optimise python code for random generator

Question

Could you please help me to optimize my first python program, I am trying to define a random generator, but somehow when I put a big number it hangs...

def rand(seed, x):
    if x == 1:
        return seed
    elif x == 2:
        return 1237
    elif x == 3:
        return 345892
    else:
        return ((31 * rand(seed, (x - 3))) + (103 * rand(seed, (x - 2))) + (7 * rand(seed, (x - 1))) + 500003) % 1000001

please try with:

print(rand(5, 5)), the result is : 506860

but when you use a bigger number its a problem, example:

print(rand(50, 50))

Answer 1

You can use dynamic programming :

def rand(seed, x):
  table = [seed, 1237, 345892]
  for _ in range(x - 3):
    next_num = table[-3] * 31 + table[-2] * 103 + table[-1] * 7 
    next_num += 500003
    next_num %= 1000001
    table.append(next_num)
  return table[-1]

The reason this is faster is that it only calculates each value once, while your code calculates things exponentially.

This algorithm is O(n) (for loop of length x - 3 ), while your algorithm was O(1.83...^n) ( 1.83^3=1+1.83+1.83^2 ).

Alternately, as suggested in the comments, you can use functools.lru_cache :

import functools

@functools.lru_cache()
def rand(seed, x):
  # insert your code here

This works in almost the same way as the previous code, but uses a decorator to memoize results instead of storing them in a list.