Ver código fonte

Clean up day 12

Taddeus Kroes 6 anos atrás
pai
commit
f9654f9326
1 arquivos alterados com 5 adições e 6 exclusões
  1. 5 6
      2019/12_jupiter.py

+ 5 - 6
2019/12_jupiter.py

@@ -27,16 +27,15 @@ def sim(axis):
 def energy_after(moons, steps):
     x, y, z = zip(*moons)
     axes = next(islice(zip(sim(x), sim(y), sim(z)), steps - 1, steps))
-    pos, vel = (tuple(zip(*i)) for i in zip(*axes))
-    return sum(abssum(p) * abssum(v) for p, v in zip(pos, vel))
+    pos, vel = zip(*axes)
+    return sum(abssum(p) * abssum(v) for p, v in zip(zip(*pos), zip(*vel)))
 
 def axis_cycle(axis):
     seen = {}
     for step, (pos, vel) in enumerate(sim(axis)):
-        ident = tuple(chain(pos, vel))
-        if ident in seen:
-            return step - seen[ident]
-        seen[ident] = step
+        prevstep = seen.setdefault(tuple(pos + vel), step)
+        if prevstep != step:
+            return step - prevstep
 
 def find_cycle(moons):
     x, y, z = zip(*moons)