Python - Memoization and Collatz Sequence

Question

When I was struggling to do Problem 14 in Project Euler, I discovered that I could use a thing called memoization to speed up my process (I let it run for a good 15 minutes, and it still hadn't returned an answer). The thing is, how do I implement it? I've tried to, but I get a keyerror(the value being returned is invalid). This bugs me because I am positive I can apply memoization to this and get this faster.

lookup = {}

def countTerms(n):
   arg = n
   count = 1
   while n is not 1:
      count += 1
      if not n%2:
         n /= 2
      else:
         n = (n*3 + 1)
      if n not in lookup:
         lookup[n] = count

   return lookup[n], arg

print max(countTerms(i) for i in range(500001, 1000000, 2)) 

Thanks.

Answer 1

There is also a nice recursive way to do this, which probably will be slower than poorsod's solution, but it is more similar to your initial code, so it may be easier for you to understand.

lookup = {}

def countTerms(n):
   if n not in lookup:
      if n == 1:
         lookup[n] = 1
      elif not n % 2:
         lookup[n] = countTerms(n / 2)[0] + 1
      else:
         lookup[n] = countTerms(n*3 + 1)[0] + 1

   return lookup[n], n

print max(countTerms(i) for i in range(500001, 1000000, 2))