Add solution for day 4

src/day04.py

@@ -0,0 +1,126 @@
+#!/usr/bin/env python3
+
+from util import print_color, read_input, Color
+
+
+def search(grid, row, col, update_row, update_col, query):
+    if query == '':
+        return True
+    if row < 0 or col < 0:
+        return False
+    try:
+        char = grid[row][col]
+        if char != query[0]:
+            return False
+    except IndexError:
+        return False
+    return search(
+        grid,
+        update_row(row),
+        update_col(col),
+        update_row,
+        update_col,
+        query[1:]
+    )
+
+
+cardinal_searches = {
+    'forward': (lambda row: row, lambda col: col + 1),
+    'backward': (lambda row: row, lambda col: col - 1),
+    'up': (lambda row: row - 1, lambda col: col),
+    'down': (lambda row: row + 1, lambda col: col),
+}
+
+
+def diagonal_searches(offset=1):
+    return {
+        'forward-up': (lambda row: row - offset, lambda col: col + offset),
+        'back-down': (lambda row: row + offset, lambda col: col - offset),
+        'forward-down': (lambda row: row + offset, lambda col: col + offset),
+        'back-up': (lambda row: row - offset, lambda col: col - offset),
+    }
+
+
+search_permutations = cardinal_searches | diagonal_searches()
+
+
+def part1(input_lines):
+    total = 0
+    for row in range(len(input_lines)):
+        for col in range(len(input_lines[row])):
+            char = input_lines[row][col]
+            if char != 'X':
+                continue
+
+            for (name, search_strategy) in search_permutations.items():
+                if search(
+                        input_lines,
+                        row,
+                        col,
+                        search_strategy[0],
+                        search_strategy[1],
+                        'XMAS'
+                ):
+                    total += 1
+    return total
+
+
+def part2(input_lines):
+    total = 0
+    x_search = diagonal_searches(2)
+    for row in range(len(input_lines)):
+        for col in range(len(input_lines[row])):
+            char = input_lines[row][col]
+            if char != 'A':
+                continue
+
+            queries = ['MS', 'SM']
+            forward_down = x_search['forward-down']
+            forward_down_matches = search(
+                input_lines,
+                row - 1,
+                col - 1,
+                forward_down[0],
+                forward_down[1],
+                queries[0]
+            ) or search(
+                input_lines,
+                row - 1,
+                col - 1,
+                forward_down[0],
+                forward_down[1],
+                queries[1]
+            )
+            forward_up = x_search['forward-up']
+            forward_up_matches = search(
+                input_lines,
+                row + 1,
+                col - 1,
+                forward_up[0],
+                forward_up[1],
+                queries[0]
+            ) or search(
+                input_lines,
+                row + 1,
+                col - 1,
+                forward_up[0],
+                forward_up[1],
+                queries[1]
+            )
+            if forward_down_matches and forward_up_matches:
+                total += 1
+    return total
+
+
+test_input = read_input("_test")
+test_input_result = part1(test_input)
+print_color(test_input_result, Color.GREY)
+assert 18 == test_input_result
+
+real_input = read_input()
+print(part1(real_input))
+
+test_input_result = part2(test_input)
+print_color(test_input_result, Color.GREY)
+assert 9 == test_input_result
+print(part2(real_input))