day 20 part 2 python

day20.py

@@ -7,14 +7,12 @@ def parse_input(input: list[str]):
     return l, w, h, start, goal
 
 
-def sp(points, w, h, start, goal):
+def sssp(points, w, h, start):
     visited = dict()
     (sx, sy) = start
     stack = [(sx, sy, 0)]
     while len(stack) > 0:
         (x, y, c) = stack.pop(0)
-        if (x, y) == goal:
-            return c
 
         if (x, y) in visited:
             continue
@@ -24,7 +22,25 @@ def sp(points, w, h, start, goal):
             if 0 <= a <= w and 0 <= b <= w and (a, b) not in points:
                 stack.append((a, b, c + 1))
 
-    return None
+    return visited
+
+
+def mhd(a, b):
+    (xa, ya) = a
+    (xb, yb) = b
+    return abs(xa - xb) + abs(ya - yb)
+
+
+def others(origin, max_distance):
+    (sx, sy) = origin
+    return [(a, b) for a in range(sx - max_distance, sx + max_distance + 1) for b in range(sy - max_distance, sy + max_distance + 1)]
+
+
+def solve(points, threshold, cheat_length):
+    fromstart = sssp(set(points), w, h, start)
+    toend = sssp(set(points), w, h, goal)
+    cheatrange = range(2, cheat_length + 1)
+    return len([1 for a in fromstart for b in others(a, cheat_length) if mhd(a, b) in cheatrange and b in toend and fromstart[a] + mhd(a, b) + toend[b] <= threshold])
 
 
 source = "input/day20.txt"
@@ -32,14 +48,10 @@ with open(source, 'r') as f:
     data = f.read()
 
 (points, w, h, start, goal) = parse_input(data.strip().split('\n'))
-bottomline = sp(set(points), w, h, start, goal)
-print(bottomline)
-atleast100 = 0
-for p in points:
-    remains = [x for x in points if p != x]
-    pathlength = sp(set(remains), w, h, start, goal)
-    if bottomline - pathlength >= 100:
-        atleast100 += 1
+bottomline = sssp(set(points), w, h, start)[goal]
 
+solution1 = solve(set(points), bottomline - 100, 2)
+print("part 1: " + str(solution1))
 
-print("task 1: " + str(atleast100))
+solution2 = solve(set(points), bottomline - 100, 20)
+print("part 2: " + str(solution2))