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()
