day 20 python cleanup

day20/sol.py

@@ -7,36 +7,34 @@ for line in open("input"):
 
 def mix(data, rounds=1):
     #assert len(set(data)) == len(data)
-    copy = list(data)
     idx = list(range(len(data)))
+    N = len(data) - 1
 
     for _ in range(rounds):
 
         for k in range(len(data)):
             i = idx.index(k)
-            x = copy[i]
+            x = data[k]
             if x != 0:
-                j = (i+x) % (len(copy)-1)
+                j = (i+x) % N
                 if x < 0 and j == 0:
-                    j = len(copy)-1
-                copy.insert(j, copy.pop(i))
+                    j = N
                 idx.insert(j, idx.pop(i))
             else:
                 j = i
-            #print(x, i, j, copy)
+            #print(x, i, j, [data[k] for k in idx])
 
-    return copy
+    return [data[k] for k in idx]
 
 def answer(data):
     i = data.index(0)
-    n = 0
+    n = []
     for j in 1000, 2000, 3000:
-        x = data[(i+j)%len(copy)]
-        #print(x)
-        n += x
-    return n
+        n.append( data[(i+j)%len(data)] )
+    print(n)
+    return sum(n)
 
-#print(copy)
+#print(data)
 
 copy = mix(data)
 print(answer(copy))