Improve lib and use it to solve 3 and 17.

d17.py

@@ -28,7 +28,7 @@ def solve(i: Input, second=False):
         pos, dirs = node
         repeats, prev_dir = dirs
         nbs = []
-        for dir in [NORTH, WEST, SOUTH, EAST]:
+        for dir in g.dirs_ort():
             if second:
                 if repeats < 4 and prev_dir is not None and prev_dir != dir:
                     continue
@@ -37,18 +37,11 @@ def solve(i: Input, second=False):
             else:
                 if repeats == 3 and prev_dir == dir:
                     continue
-
+            if prev_dir == g.flip_ort(dir):
-            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):
+            if nb not in g:
                 continue
             nbs.append((nb, (repeats + 1 if dir == prev_dir else 1, dir)))
         return nbs

d3.py

@@ -1,6 +1,6 @@
 import re
+from lib import *
 
-NUMBERS = "0123456789"
 EXAMPLE = """
 467..114..
 ...*......
@@ -14,6 +14,42 @@ EXAMPLE = """
 .664.598..
 """
 
+def solve(i: Input, second=False):
+    """
+    This is a new version I have implemented after the fact to test the
+    improvements to my library. My original versions are `solve1` and `solve2`.
+    """
+    g = i.grid2()
+    if not second:
+        parts = g.find_not(NUMBERS + ".")
+    else:
+        parts = g.find("*")
+
+    res = 0
+    numbers = []
+    for p in parts:
+        numbers_gear = []
+        for n in g.neighbors_adj(p):
+            if g[n] in NUMBERS:
+                while n in g and g[n] in NUMBERS:
+                    n = (n[0], n[1] - 1)
+                number = ""
+                n = (n[0], n[1] + 1)
+                start = n
+                while n in g and g[n] in NUMBERS:
+                    number += g[n]
+                    n = (n[0], n[1] + 1)
+                numbers.append((int(number), n[0], start))
+                numbers_gear.append(int(number))
+        numbers_gear = list(set(numbers_gear))
+        if len(numbers_gear) == 2:
+            res += numbers_gear[0] * numbers_gear[1]
+
+    if second:
+        return res
+    else:
+        return sum([n for n, _, _ in list(set(numbers))])
+
 def clean(text: str) -> list[str]:
     return list(filter(lambda l: l.strip() != "", text.splitlines()))
 
@@ -29,7 +65,7 @@ def is_adj_to_symbol(x, y, lines):
             pass
     return False
 
-def solve(lines: list[str]):
+def solve1(lines: list[str]):
     d = ""
     adj_to_symbol = False
     s = 0
@@ -86,16 +122,18 @@ def solve2(lines: list[str]):
     return s
 
 def main():
-    example = clean(EXAMPLE)
+    input = Input(EXAMPLE)
-    print("Example  1:", solve(example))
+    print("Example  1:", solve(input))
-    data = clean(open("i3.txt").read())
-    print("Solution 1:", solve(data))
 
-    example = clean(EXAMPLE)
+    input = Input("i3.txt")
-    print("Example  2:", solve2(example))
+    print("Solution 1:", solve(input))
 
-    data = clean(open("i3.txt").read())
+    input = Input(EXAMPLE)
-    print("Solution 2:", solve2(data))
+    print("Example  2:", solve(input, True))
+
+    input = Input("i3.txt")
+    print("Solution 2:", solve(input, True))
+    return
 
 if __name__ == "__main__":
     main()

lib.py

@@ -7,16 +7,6 @@ NUMBERS = string.digits
 LETTERS_LOWER = string.ascii_lowercase
 LETTERS_UPPER = string.ascii_uppercase
 
-UP = (-1, 0)
-DOWN = (1, 0)
-RIGHT = (0, 1)
-LEFT = (0, -1)
-
-NORTH = UP
-SOUTH = DOWN
-EAST = RIGHT
-WEST = LEFT
-
 INF = float("inf")
 fst = lambda l: l[0]
 snd = lambda l: l[1]
@@ -33,6 +23,15 @@ def add2(a: tuple[int, int], b: tuple[int, int]) -> tuple[int, int]:
     return (a[0] + b[0], a[1] + b[1])
 
 class Grid2D:
+    N = (-1, 0)
+    E = (0, 1)
+    S = (1, 0)
+    W = (0, -1)
+    NW = (-1, -1)
+    NE = (-1, 1)
+    SE = (1, 1)
+    SW = (1, -1)
+
     def __init__(self, text: str):
         lines = [line for line in text.splitlines() if line.strip() != ""]
         self.grid = list(map(list, lines))
@@ -64,13 +63,10 @@ class Grid2D:
                 for col_i in range(self.n_cols)]
 
     def find(self, chars: str) -> list[tuple[int, int]]:
-        r = []
+        return [c for c in self.all_coords() if self[c] in chars]
-        for row_i in range(self.n_rows):
+
-            for col_i in range(self.n_cols):
+    def find_not(self, chars: str) -> list[tuple[int, int]]:
-                c = (row_i, col_i)
+        return [c for c in self.all_coords() if self[c] not in chars]
-                if self[c] in chars:
-                    r.append(c)
-        return r
 
     def all_coords(self) -> list[tuple[int, int]]:
         return [(row_i, col_i)
@@ -89,17 +85,27 @@ class Grid2D:
         row, col = pos
         return row >= 0 and row < self.n_rows and col >= 0 and col < self.n_cols
 
+    def __contains__(self, pos: tuple[int, int]) -> bool:
+        return self.contains(pos)
+
     def neighbors_ort(self, pos: tuple[int, int]) -> list[tuple[int, int]]:
-        ort_rel = [(-1, 0), (0, 1), (1, 0), (0, -1)]
+        return [add2(pos, off) for off in self.dirs_ort() if self.contains(add2(pos, off))]
-        return [add2(pos, off) for off in ort_rel if self.contains(add2(pos, off))]
 
     def neighbors_vert(self, pos: tuple[int, int]) -> list[tuple[int, int]]:
-        ort_vert = [(-1, -1), (-1, 1), (1, 1), (1, -1)]
+        return [add2(pos, off) for off in self.dirs_vert() if self.contains(add2(pos, off))]
-        return [add2(pos, off) for off in ort_vert if self.contains(add2(pos, off))]
 
     def neighbors_adj(self, pos: tuple[int, int]) -> list[tuple[int, int]]:
         return self.neighbors_ort(pos) + self.neighbors_vert(pos)
 
+    def flip_ort(self, pos: tuple[int, int]) -> tuple[int, int]:
+        return (-pos[0], -pos[1])
+
+    def dirs_ort(self) -> list[tuple[int, int]]:
+        return [self.N, self.E, self.S, self.W]
+
+    def dirs_vert(self) -> list[tuple[int, int]]:
+        return [self.NE, self.SE, self.SW, self.NW]
+
     def print(self):
         for r in self.rows():
             print("".join(r))