Extend lib and solve day 17 with A*. Fun.

d17.py

@@ -0,0 +1,78 @@
+from lib import *
+
+EXAMPLE = """
+2413432311323
+3215453535623
+3255245654254
+3446585845452
+4546657867536
+1438598798454
+4457876987766
+3637877979653
+4654967986887
+4564679986453
+1224686865563
+2546548887735
+4322674655533
+"""
+
+def solve(i: Input, second=False):
+    g = i.grid2()
+    starts = [((0, 0), (0, None))]
+
+    def is_goal(node):
+        pos, _ = node
+        return pos == (g.n_cols - 1, g.n_rows - 1)
+
+    def neighbors(node):
+        pos, dirs = node
+        repeats, prev_dir = dirs
+        nbs = []
+        for dir in [NORTH, WEST, SOUTH, EAST]:
+            if second:
+                if repeats < 4 and prev_dir is not None and prev_dir != dir:
+                    continue
+                if repeats == 10 and prev_dir == dir:
+                    continue
+            else:
+                if repeats == 3 and prev_dir == dir:
+                    continue
+
+            if prev_dir == NORTH and dir == SOUTH:
+                continue
+            elif prev_dir == SOUTH and dir == NORTH:
+                continue
+            elif prev_dir == EAST and dir == WEST:
+                continue
+            elif prev_dir == WEST and dir == EAST:
+                continue
+
+            nb = add2(pos, dir)
+            if not g.contains(nb):
+                continue
+            nbs.append((nb, (repeats + 1 if dir == prev_dir else 1, dir)))
+        return nbs
+
+    def h(node):
+        pos, _ = node
+        return abs(g.n_rows - 1 - pos[0]) + abs(g.n_cols - 1 - pos[1])
+
+    def d(_, b):
+        pos, _ = b
+        if pos == (0, 0):
+            return 0
+        return int(g[pos])
+
+    a = A_Star(starts, is_goal, h, d, neighbors)
+    return a.cost
+
+def main():
+    DAY_INPUT = "i17.txt"
+    print("Example  1:", solve(Input(EXAMPLE)))
+    print("Solution 1:", solve(Input(DAY_INPUT)))
+    print("Example  2:", solve(Input(EXAMPLE), True))
+    print("Solution 2:", solve(Input(DAY_INPUT), True))
+    return
+
+if __name__ == "__main__":
+    main()