day20.py (1591B)
1 import sys 2 import numpy as np 3 4 sys.setrecursionlimit(15000) 5 6 saved_u = None 7 8 def u(i, j, s): 9 if saved_u[i,j] != 0xffff: 10 return saved_u[i,j] 11 (n, d, zero_index) = s 12 assert(i >= 0) 13 assert(j >= 0) 14 assert(i <= n) 15 assert(j < n) 16 if i == 0: 17 if j >= zero_index: 18 r = j-zero_index 19 else: 20 r = j+n-zero_index 21 else: 22 ti = (u(i-1, i-1, s) + d[i-1]) % (n-1) 23 if i-1 == j: 24 print("i = {}, ti = {}".format(i, ti)) 25 if d[i-1] != 0 and ti == 0: 26 tip = n-1 27 else: 28 tip = ti 29 if i-1 == j: 30 r = tip 31 elif tip <= u(i-1, j, s) < u(i-1, i-1, s): 32 r = u(i-1, j, s) + 1 33 elif u(i-1, i-1, s) < u(i-1, j, s) <= tip: 34 r = u(i-1, j, s) - 1 35 else: 36 r = u(i-1, j, s) 37 saved_u[i, j] = r 38 return r 39 40 41 d = [] 42 f = open("input20", "r") 43 for line in f.readlines(): 44 if line[-1] == "\n": 45 line = line[:-1] 46 d.append(int(line)) 47 print(d) 48 n = len(d) 49 its = 47 50 saved_u = np.array([[0xffff]*(n+1)]*(n+1), np.ushort) 51 print(saved_u) 52 zero_index = d.index(0) 53 s = (n, d, zero_index) 54 print("n = {}, zero_index = {}".format(n, zero_index)) 55 56 for i in range(its): 57 for j in range(n): 58 up = u(i,j,s) 59 us = [u(i, j, s) for j in range(n)] 60 js = [x % n for x in [1000, 2000, 3000]] 61 su = sum([d[us.index(j)] for j in js]) 62 print(su) 63 # print([u(i, j, s) for j in range(n)]) 64 65 #us = [u(its, j, s) for j in range(n)] 66 #js = [x % n for x in [1000, 2000, 3000]] 67 #print(sum([d[us.index(j)] for j in js]))