4
\$\begingroup\$

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.

\$\endgroup\$
  • \$\begingroup\$ Can you add a main function demonstrating what you program actually does? It is for computing the Collatz sequence for a single number? Or does it compute many Collatz sequences? Or does it try to find the longest sequence (in some range)? \$\endgroup\$ – Martin R Nov 4 at 15:39
  • \$\begingroup\$ Any/all of the above really. I want to be efficient in time and space, accounting for repeated invocations regardless of same/new inputs. I'll edit it into the question. \$\endgroup\$ – Gloweye Nov 4 at 15:45
5
\$\begingroup\$

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.

\$\endgroup\$
  • \$\begingroup\$ Looks good. that's the sort of caching I was looking for, I guess. I stared at the code for 3 hours without finding a neat way to decorator-cache it like that. I'm letting others have a bit of a shot till sometime tomorrow, but unless I get a good one you're gonna be accepted. \$\endgroup\$ – Gloweye Nov 4 at 18:06
  • \$\begingroup\$ @Gloweye I'm not sure where you got to while implementing your own decorator, but I added a minimal reference-implementation in case you wanted to see how it could be written. Far from perfect, but it gets the job done. \$\endgroup\$ – Carcigenicate Nov 4 at 18:22
  • \$\begingroup\$ Nah, more like that I never found a nice way to factor it into a function to decorate. Probably stared myself blind on the code. \$\endgroup\$ – Gloweye Nov 4 at 18:41
  • \$\begingroup\$ This is all true, but you don't use the fact that once you hit the cache, you know for sure all the subsequent values will also be in the cache. \$\endgroup\$ – One Lyner Nov 6 at 14:46
  • 1
    \$\begingroup\$ @OneLyner Although I suppose looking over your version, it isn't any more complicated, just moved around. \$\endgroup\$ – Carcigenicate Nov 6 at 14:57
3
\$\begingroup\$

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.

\$\endgroup\$
1
\$\begingroup\$

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)
\$\endgroup\$
  • \$\begingroup\$ in quick and dirty (manual) timing: If I run [collatz(n) for n in range(1, 1000000)] for mine, it's done in 20 seconds, while with your's it's about 28 seconds. And that's still considering mine needs to generate a full list for every sequence while yours doesn't bother with that. But the idea was mostly about reviewing my actual code, and pointing out which parts of my code suck ;) \$\endgroup\$ – Gloweye Nov 4 at 12:55
  • 1
    \$\begingroup\$ The "complex logic" saves you from recalculating each step if it had be ccukated previously. That can save, at best, calls to %, which may be faster. Memorization/caching is useful when you need to carry out an expensive operation over and over, and can expect results to repeat themselves. \$\endgroup\$ – Carcigenicate Nov 6 at 14:58
  • \$\begingroup\$ thanks @Carcigenicate, this is what I was looking for. \$\endgroup\$ – Digvijay Rawat Nov 6 at 15:36

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.