Calculating a Collatz Sequence

Question

I've written a small script to calculate Collatz sequences. For background, read here about the Collatz conjecture. Largest number I've tried so far in my implementation is 93571393692802302 (which wikipedia claims is the number requiring most steps below 10e17), which needs 2091 steps. Apparently I count different from Wikipedia, as my len(collatz(93571393692802302)) is 2092 long, but I include the number itself as well as the final 1 .

A Collatz sequence takes a number, and calculates the next number in sequence as following:

• If it is even, divide by two
• If it is odd, triple and add one.

The Collatz Conjecture claims that all sequences created this way converge to 1. It is unproven, but it seems most mathematicians suspect it's True.

Without further ado, the code:

from typing import List


collatz_cache = {}


def _calc(number: int) -> int:
    """
    Caches this number in the collatz_cache, and returns the next in sequence.
    :param number: Integer > 1.
    :return: Integer > 0
    """
    if number % 2 == 0:
        next_number = number // 2
    else:
        next_number = number * 3 + 1
    collatz_cache[number] = next_number
    return next_number


def collatz(number: int) -> List[int]:
    """
    Returns the collatz sequence of a number.
    :param number: Integer > 0
    :return: Collatz sequence starting with number
    """
    if not isinstance(number, int):
        raise TypeError(f"Collatz sequence doesn't make sense for non-integers like {type(number).__name__}")
    if number < 1:
        raise ValueError(f"Collatz sequence not defined for {type(number).__name__}({number})")
    new_number = number
    result = [number]
    while new_number not in collatz_cache and new_number != 1:
        new_number = _calc(new_number)
        result.append(new_number)
    while result[-1] != 1:
        result.append(collatz_cache[result[-1]])
    return result

I've tried to minimize calculation time in repeated attempts by creating a mapping for each number to the next number in the sequence. First I just mapped to the full list, but that would make the dictionary a bit bigger than I really wanted.

I feel that there should be some gains to be made in the list construction, but I'm at loss as to the how.

Purpose of this code is to basically be a general library function. I want to:

• Be fast
• Be memory efficient with my cache
• Handle multiple equal/partially overlapping sequences
• Handle completely different sequences

And all of that at the same time. Any code stylistic improvements are of course welcome, but any suggestions improving the above goals without disproportional disadvantage to the others are also welcome.

Answer 1

I think the handling of the cache could be improved. You add to the cache in _calc, then check the cache in collatz. Why not just let _calc handle the cache? That way collatz doesn't need to know anything about how _calc is getting its results. I'd change it to something like:

def _calc2(number: int) -> int:
    if number in collatz_cache:  # Let it do the lookup
        return collatz_cache[number]

    if number % 2 == 0:
        next_number = number // 2
    else:
        next_number = number * 3 + 1

    collatz_cache[number] = next_number
    return next_number

def collatz2(number: int) -> List[int]:
    if number < 1:
        raise ValueError(f"Collatz sequence not defined for {type(number).__name__}({number})")

    results = [number]

    new_number = number
    while new_number != 1:
        new_number = _calc2(new_number)
        results.append(new_number)

    return results

I also don't think collatz really needs to do an isinstance check on number. You've said via type-hints that it only accepts integers. If you're going to go down that road, arguably every function should check what it's arguments are. I suppose protecting against a float being passed may be worth it since that will silently mess with the results, but you can only hand-hold the user so much.

I'll note though that your manual memoization isn't necessary. functools has a decorator for this:

from functools import lru_cache

@lru_cache(maxsize=int(1e6))  # This could probably be smaller
def _calc3(number: int) -> int:
    if number % 2 == 0:
        next_number = number // 2
    else:
        next_number = number * 3 + 1

    return next_number

Using a dirty timeit timing using seeded random numbers, I found that all three versions perform nearly identically for me:

from timeit import timeit
from random import randint, seed

timeit(lambda: collatz(randint(5, 100)),
       setup=lambda: seed(12345),
       number=int(1e6))

# All take between 7 seconds for me

I used random numbers so the cache is actually getting fully tested instead of just returning the same numbers over and over. The random numbers are seeded, so the results should be reliable.

I'll note, you could also write your own bare-bones memoize decorator as well:

def memoize(f):
    cache = {}

    def wrapper(*args):
        if args in cache:
            return cache[args]

        else:
            result = f(*args)
            cache[args] = result
            return result

    return wrapper

@memoize
def tester(n, m):
    print("Called with", n, m)
    return n + m

>>>> tester(1, 3)
Called with 1 3
4

>>> tester(1, 3)
4

---

I'll note too though, this could be made into a generator if you wanted to generalize it:

from typing import Generator

def collatz3(number: int) -> Generator[int, None, None]:
    yield number  # To be consistent with the other functions

    new_number = number
    while new_number != 1:
        new_number = _calc3(new_number)
        yield new_number

This may prove to be more memory efficient for longer sequences, depending on how it's used. Now the whole sequence isn't required to be held in memory all at once.

Answer 2

I agree with @Carcigenicate that splitting the cache responsibility between two functions is not a good idea. But contrary to him, I would make the main collatz function handle it, because it has the specific knowledge that once you hit the cache, you know the rest of the sequence will also be in there.

Here is an implementation.

from typing import List

_cache = { 1: None }

def _next(i: int) -> int:
    """
    Returns the next element in the Collatz sequence
    """
    if i % 2:
        return 3 * i + 1
    else:
        return i // 2

def collatz(i: int) -> List[int]:
    """
    Returns the collatz sequence of a number.
    :param number: Integer > 0
    :return: Collatz sequence starting with number
    """
    r=[] 
    while i not in _cache: 
        r.append(i) 
        n = _next(i) 
        _cache[i] = n 
        i = n
    while i: 
        r.append(i) 
        i = _cache[i] 
    return r

This makes this version a bit faster for me.

Answer 3

I am self learning Python and am not so experienced as you are. So I actually have a couple of questions on your code. Mainly, I want to know why have you built such a complex logic for such a simple problem? TBH, I am not even sure if I understand fully what your code is doing. Below I have pasted my code, and it gives the right answer (2091 steps) in 0.0005s, which is small enough I guess. Are you trying to run it in lesser time duration than that, and hence the complicated logic?

def collatz(n):
    steps = 0
    while n != 1:
        if n % 2 == 0:
            n = n // 2
        else:
            n = (n * 3) + 1
        steps += 1

    print(steps)