10
\$\begingroup\$
def collatz(n, counter):
    if n == 1:
        return counter
    elif n % 2 == 0:
        return collatz(n/2, counter + 1)
    else:
        return collatz(3*n + 1, counter + 1)

print(int(collatz(15, 0)))

Is there any way to improve this code? The two arguments passed on the collatz() function rubs me the wrong way. Tells me that should only be one, but I don't know better.

My question is an attempt to improving my Python vocabulary as I just barely started. So, what useful Python tool could I have used here to make the code better?

\$\endgroup\$
5
  • \$\begingroup\$ Do Python implementations typically eliminate tail-calls? If not, you're heading for a stack overflow by using recursion... \$\endgroup\$ Aug 26, 2021 at 7:38
  • 4
    \$\begingroup\$ @TobySpeight CPython, the reference implementation, doesn't perform any form of tail call elimination \$\endgroup\$
    – Jasmijn
    Aug 26, 2021 at 15:21
  • 2
    \$\begingroup\$ See tobiaskohn.ch/index.php/2018/08/28/optimising-python-3 for more information @TobySpeight \$\endgroup\$ Aug 26, 2021 at 16:23
  • 3
    \$\begingroup\$ @Jasmijn: It's not just CPython. Guido van Rossum has said that NO Python implementation is allowed to eliminate tail calls. So, any implementation that eliminates tail calls is non-compliant, and thus by definition not a Python implementation. Therefore, it is impossible that there is a Python implementation that eliminates tail calls, because if it did, it wouldn't be a Python implementation. (Personally, I find that quite insane: not eliminating tail calls is essentially a memory leak, so why would you force implementors to leak memory?) \$\endgroup\$ Aug 26, 2021 at 22:47
  • \$\begingroup\$ @JörgWMittag Where do you store traceback information if TCO is allowed? In another stack? Isn't the point that tracebacks are more important than allowing you to use recursion to loop? If you need such levels of recursion you can normally quite easily convert away from an FP approach. \$\endgroup\$
    – Peilonrayz
    Aug 27, 2021 at 9:37

2 Answers 2

9
\$\begingroup\$

Formatting / Spacing

There should be spaces around operators, such as n / 2 or 3 * n. Some IDEs can handle this for you through an auto-formatting option (e.g. Ctrl + Alt + L for PyCharm on Windows).


Type hints

It's useful to provide type hints for arguments and function return values. This increases readability and allows for easier and better error checking. Refer to PEP484 for more information.

def collatz(n: int, counter: int) -> int:
    # function body here

Return values

Adding these type hints will help us identify the first improvement. As the Collatz sequence only contains integers, our collatz function should only take an integer as the n argument. collatz(n / 2, counter + 1) passes a float, so to keep it consistent we should probably convert it to an int before passing it: collatz(int(n / 2), counter + 1). Even better, we can use the floor division operator //: collatz(n // 2, counter + 1)

Please note that you do not need to convert the function output to an int, as it will only ever return an int value.


Default arguments

There are multiple approaches to improving the handling of the counter argument, which rubs you the wrong way. This is good intuition on your part. With default arguments, Python offers one really concise option:

def collatz(n: int, counter: int = 0) -> int:
    if n == 1:
        return counter
    elif n % 2 == 0:
        return collatz(n // 2, counter + 1)
    else:
        return collatz(3 * n + 1, counter + 1)

As you can see, the only addition we need to implement is counter = 0, which makes 0 the default argument for counter. This means that counter will be set to 0 if the argument is not provided by the caller.

You can now simply call

print(collatz(15))

More on default arguments in Python.

\$\endgroup\$
11
  • 5
    \$\begingroup\$ I would also recommend using cache from itertools, it will dramatically speed up the code. \$\endgroup\$ Aug 26, 2021 at 7:58
  • \$\begingroup\$ @N3buchadnezzar I agree, good recommendation! \$\endgroup\$ Aug 26, 2021 at 8:01
  • 2
    \$\begingroup\$ The type hints are certainly useful, and I definitely miss it coming from C. Do Python programmers regularly use them? Or is it something that goes away when you're already comfortable using Python? The counter variable initialization was my problem; I could not find where it should be put. And so my solution was just to pass it to the function. I know that's not how it should be because I admit my Python knowledge is still very limited. You provided me with easy-to-understand explanations. Thank you :) \$\endgroup\$ Aug 26, 2021 at 8:01
  • 1
    \$\begingroup\$ Glad to help! I'd say type hints are useful at any level of Python proficiency, as they allow for easier debugging and maintainability. They do probably provide the biggest benefit to beginner and intermediate Pythonistas though =) \$\endgroup\$ Aug 26, 2021 at 8:34
  • \$\begingroup\$ Btw why do we care about keeping track of counter? \$\endgroup\$ Aug 26, 2021 at 8:48
15
\$\begingroup\$

Sequence vs counting

As stated in the comments

the goal of the code is to print the length of the Collatz sequence. Could you elaborate as to why you asked?

As OP mentions, he is not interested in the sequence itself, but its length. As such we actually do not need the values from the sequence itself. We only need to count how many iterations it takes to reach 1. The following code does precisely that: every time the function is called we increment by one:

def collatz(n: int) -> int:
    if n == 1:
        return 1
    elif n % 2 == 0:
        return 1 + collatz(n // 2)
    else:  # n % 2 == 1:
        return 1 + collatz(3 * n + 1)

Spend some time thinking about this. Recursion is hard. Go through the code above by hand for the number 5 and see what it returns and how. As a minor point, it is better to be explicit than implicit in Python. Compare

if n == 1:
    return n

vs

if n == 1:
    return 1

While trivial, it is a good mindset to get into.

Cache

It can be very wise to cache previous calls to the function to save time. Assume we try to calculate collatz(23):

23, 70, 35, 106, 53, 160, 80, 40, 20, 10, 5, 16, 8, 4, 2, 1 

So collatz(23) = 15. Now assume we want to calculate collatz(61):

 61, 184, 92, 46, (23)

Notice how we stop early: 23 is already saved so we only have to do 4 iterations instead of 19. This can, for instance, be implemented as follows:

cache = {1: 0}

def collatz(n: int) -> int:
    if n in cache:
        return cache[n]
    else:
        if n % 2 == 0:
            m = n // 2
        else:
            m = 3 * n + 1
        res = collatz(m) + 1
        cache[n] = res
        return res

However. there are builtins for handling memoization in Python. Introducing the decorator itertools.cache.

import functools


@functools.cache
def collatz(n: int) -> int:
    if n == 1:
        return 1
    elif n % 2 == 0:
        return 1 + collatz(n // 2)
    else:  # n % 2 == 1:
        return 1 + collatz(3 * n + 1)

Let us add a test function to benchmark how much quicker our function is with memoization:

def longest_collatz(limit: int) -> int:
    longest = 0
    for i in range(1, limit):
        current = collatz(i)
        if current > longest:
            longest = current
    return longest


def main():
    limit = 10 ** 4
    with cProfile.Profile() as pr:
        longest_collatz(limit)

    stats = pstats.Stats(pr)
    stats.strip_dirs()
    stats.sort_stats(pstats.SortKey.CALLS)
    stats.print_stats()

Here we simply compare how many function calls it takes to find the longest Collatz sequence amongst the first 10 000 numbers. I wanted to try with higher values but your version took too long to complete...

859639 function calls (10002 primitive calls) in 12.444 seconds
 21667 function calls ( 4330 primitive calls) in  0.332 seconds

Of course it is much smarter to just iterate over the odd values, but this is just for comparison. To compare the versions I just commented the @functools.cache bit.

Full code

import functools
import cProfile
import pstats

@functools.cache
def collatz(n: int) -> int:
    if n == 1:
        return n
    elif n % 2 == 0:
        return 1 + collatz(n // 2)
    else:  # n % 2 == 1:
        return 1 + collatz(3 * n + 1)

def longest_collatz(limit: int) -> int:
    longest = 0
    for i in range(1, limit):
        current = collatz(i)
        if current > longest:
            longest = current
    return longest

def main():
    limit = 10 ** 4
    with cProfile.Profile() as pr:
        longest_collatz(limit)

    stats = pstats.Stats(pr)
    stats.strip_dirs()
    stats.sort_stats(pstats.SortKey.CALLS)
    stats.print_stats()

if __name__ == "__main__":
    main()
\$\endgroup\$
3
  • 1
    \$\begingroup\$ I could imagine it would be an awesome feeling to go back to this answer when I can finally understand every concept you put here. I highly appreciate how you optimized the code and explained it in a way that's not daunting. The runtime data is the icing on the cake. Thank you :) \$\endgroup\$ Aug 26, 2021 at 15:23
  • \$\begingroup\$ Note that we could have done this even faster if our goal was not to generate the length of the sequence, but the longest one. I wrote some bad (but super fast) code for solving this a bunch of years ago github.com/Oisov/Project-Euler/blob/master/Problems/PE_014/…. In python 2, but it can easily be rewritten. The next step up after understanding this answer, is understanding this code. The final step is improving the code with typing hints, doctests and proper docstring + comments. \$\endgroup\$ Aug 26, 2021 at 16:21
  • 1
    \$\begingroup\$ As a final comment the github link also avoids the tail call recursion issue by simply implementing a iterative instead of recursive version =) \$\endgroup\$ Aug 26, 2021 at 16:28

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.