day 9 python simplify

day09/day09.py

@@ -1,84 +1,42 @@
-H = (0,0)
-T = (0,0)
 def go(p, dir):
     x,y = p
-    if dir == "U":
-        y += 1
-    if dir == "D":
-        y -= 1
-    if dir == "L":
-        x -= 1
-    if dir == "R":
-        x += 1
+    if   dir == "U": y += 1
+    elif dir == "D": y -= 1
+    elif dir == "L": x -= 1
+    elif dir == "R": x += 1
     return x,y
 
-def follow():
-    if H == T:
-        return T
-    if abs(H[0] - T[0]) <= 1 and abs(H[1] - T[1]) <= 1:
-        return T
-    if H[1] == T[1]:
-        x,y = T
-        x = H[0] - sgn(H[0] - T[0])
-        return x,y
-    if H[0] == T[0]:
-        x,y = T
-        y = H[1] -  sgn(H[1]-T[1])
-        return x,y
-    # diagonal
-    x,y = T
-    x += sgn(H[0] - T[0])
-    y += sgn(H[1] - T[1])
-    return x,y
-
-def follow2(h, t):
-    if h == t:
-        return t
-    if abs(h[0] - t[0]) <= 1 and abs(h[1] - t[1]) <= 1:
-        return t
-    if h[1] == t[1]:
+def follow(h, t):
+    if abs(h[0] - t[0]) > 1 or abs(h[1] - t[1]) > 1:
         x,y = t
-        x = h[0] - sgn(h[0] - t[0])
-        return x,y
-    if h[0] == t[0]:
-        x,y = t
-        y = h[1] -  sgn(h[1]-t[1])
-        return x,y
-    # diagonal
-    x,y = t
-    x += sgn(h[0] - t[0])
-    y += sgn(h[1] - t[1])
-    return x,y
+        x += sgn(h[0] - t[0])
+        y += sgn(h[1] - t[1])
+        t = x,y
+    return t
 
 def sgn(x):
-    if x==0: return 0
-    if x > 0: return 1
-    return -1
+    return (x > 0) - (x < 0)
 
-visited = set([T])
-for line in open("input"):
-    dir, count = line.split()
-    count = int(count)
-    for _ in range(count):
-        H = go(H, dir)
-        T = follow()
-        visited.add(T)
-        #print(H,T)
-print(len(visited))
-        
+def solve(n):
+    R = [(0,0)]*n
+    visited = set([R[-1]])
+    for line in open("input"):
+        dir, count = line.split()
+        for _ in range(int(count)):
+            R[0] = go(R[0], dir)
+            for i in range(1, len(R)):
+                t = follow(R[i-1], R[i])
+                if t == R[i]:
+                    # optimization: if a knot doesn't move,
+                    # none of the knots after it will either
+                    break
+                R[i] = t
+            visited.add(R[-1])
+            #print(R)
+    return len(visited)
 
-R = [(0,0)]*10
-visited = set([R[-1]])
-for line in open("input"):
-    dir, count = line.split()
-    count = int(count)
-    for _ in range(count):
-        R[0] = go(R[0], dir)
-        for i in range(1,len(R)):
-            t = follow2(R[i-1], R[i])
-            if t == R[i]:
-                break
-            R[i] = t
-        visited.add(R[-1])
-        print(R)
-print(len(visited))
+A1 = solve(2)
+print(A1)
+
+A2 = solve(10)
+print(A2)