2
\$\begingroup\$

I created this program originally for Project Euler 14.

Here's the code:

from sys import setrecursionlimit
setrecursionlimit(10 ** 9)

def memo(f):
    f.cache = {}

    def _f(*args):
        if args not in f.cache:
            f.cache[args] = f(*args)
        return f.cache[args]

    return _f

@memo
def collatz(n):
    if n == 1:
        return 1

    if n % 2:
        print(3 * n + 1)
        return 1 + collatz(3 * n + 1)

    if not n % 2:
        print(n // 2)
        return 1 + collatz(n // 2)

if __name__ == '__main__':
    number = None

    while True:
        try:
            number = input('Please enter a number: ')
            number = int(number)
            break

        except ValueError:
            print(f'Value has to be an interger instead of "{number}"\n')

    print('\nThe collatz sequence: ')

    length = collatz(number)

    print(f'\nLength of the collatz sequence for {number}: {length}')

I would like to make this faster and neater if possible.

\$\endgroup\$

1 Answer 1

2
\$\begingroup\$

Useless Code

Your @memo-ization doesn't do anything of value:

  • Your mainline calls collatz() exactly once. Unless given the value 1, the collatz() function calls itself with 3*n + 1 or n // 2, and since the Collatz sequence doesn't have any loops until the value 1 is reached, it will never call itself with a value it has already been called with. So, you are memoizing values which will never, ever be used.
  • Your collatz() function doesn't just return a value; it has a side-effect: printing! If you did call collatz() with a value that has already been memoized, nothing will be printed.

Tail Call, without Tail Call Optimization

You've increased the recursion limit to \$10^9\$ stack levels, which is impressive. But the algorithm doesn't need to be recursive. A simple loop would work. And since Python cannot do Tail Call Optimization, you should replace recursion with loops wherever possible. Doing so eliminates the need for the increased stack limit:

def collatz(n):
    count = 1

    while n > 1:
        if n % 2 != 0:
            n = n * 3 + 1
        else:
            n //= 2

        print(n)
        count += 1

    return count

Separate Sequence Generation from Printing

You can create a very simple Collatz sequence generator:

def collatz(n):
    while n > 1:
        yield n
        n = n * 3 + 1 if n % 2 else n // 2
    yield 1

Using this generator, the caller can print out the Collatz sequence with a simple loop:

for x in collatz(13):
    print(x)

Or, if the caller just wants the length of the Collatz sequence, without printing out each item in the sequence, you can determine the length of the sequence with len(list(collatz(13))). Better would be to count the items returned by the generator without realizing the list in memory: sum(1 for _ in collatz(13)).

Project Euler 14

The above generator works great for determining the sequence from any arbitrary value. If you want to compute the length of collatz(n) for \$1 \le n \le 1,000,000\$, you may want to return to memoization. However, this is Code Review, and while you alluded to PE14, you didn't actually provide code for that, so that cannot be reviewed here.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

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