optim.py (1029B) [raw]
1 #!/usr/bin/env python3 2 import sys 3 import os 4 5 def parse(): 6 return [int(i) for i in next(sys.stdin).split()] 7 8 def debug(msg): 9 if os.getenv("DEBUG"): 10 print(msg) 11 12 # f(x, n) = count of nodes after iterating n steps 13 # 14 # f(x, n) is either 15 # f(1, n-1) if x == 0 16 # f(y, n-1) + f(z, n-1) if x even dig 17 # f(2024*x, n-1) otherwise 18 # 19 # We can memoize the output, since eventually it will cycle! 20 cache = {} 21 def count(stone, n): 22 if n == 0: 23 return 1 24 if (stone, n) in cache: 25 return cache[(stone, n)] 26 if stone == 0: 27 c = count(1, n - 1) 28 else: 29 ss = str(stone) 30 mid = len(ss) // 2 31 l, r = ss[0:mid], ss[mid:] 32 if len(l) == len(r): 33 c = count(int(l), n-1) + count(int(r), n-1) 34 else: 35 c = count(stone * 2024, n-1) 36 cache[(stone, n)] = c 37 return c 38 39 if __name__ == '__main__': 40 stones = parse() 41 tot = 0 42 n = int(os.getenv("N", 25)) 43 for s in stones: 44 tot += count(s, n) 45 print(tot)