Implement 2024 day 17 part 2

Author: bertptrs

2024/src/aoc/days/day17.py

@@ -3,6 +3,54 @@
 from . import SeparateRunner
 
 
+def run_program(
+    register_a: int, register_b: int, register_c: int, program: list[int]
+) -> list[int]:
+    ip = 0
+    out = []
+
+    def combo(index: int) -> int:
+        match program[index]:
+            case 0:
+                return 0
+            case 1:
+                return 1
+            case 2:
+                return 2
+            case 3:
+                return 3
+            case 4:
+                return register_a
+            case 5:
+                return register_b
+            case 6:
+                return register_c
+
+    while ip < len(program):
+        match program[ip]:
+            case 0:  # adv
+                register_a = register_a // 2 ** combo(ip + 1)
+            case 1:  # bxl
+                register_b ^= program[ip + 1]
+            case 2:  # bst
+                register_b = combo(ip + 1) & 0x7
+            case 3:  # jnz
+                if register_a != 0:
+                    ip = program[ip + 1]
+                    continue
+            case 4:  # bxc
+                register_b ^= register_c
+            case 5:  # out
+                out.append(combo(ip + 1) & 7)
+            case 6:  # bdv
+                register_b = register_a // 2 ** combo(ip + 1)
+            case 7:  # cdv
+                register_c = register_a // 2 ** combo(ip + 1)
+        ip += 2
+
+    return out
+
+
 class DayRunner(SeparateRunner):
     @classmethod
     def part1(cls, input: str) -> str:
@@ -11,50 +59,31 @@ def part1(cls, input: str) -> str:
         register_a, register_b, register_c = map(int, numbers[:3])
         program = list(map(int, numbers[3:]))
 
-        ip = 0
-        out = []
-
-        def combo(index: int) -> int:
-            match program[index]:
-                case 0:
-                    return 0
-                case 1:
-                    return 1
-                case 2:
-                    return 2
-                case 3:
-                    return 3
-                case 4:
-                    return register_a
-                case 5:
-                    return register_b
-                case 6:
-                    return register_c
-
-        while ip < len(program):
-            match program[ip]:
-                case 0:  # adv
-                    register_a = register_a // 2 ** combo(ip + 1)
-                case 1:  # bxl
-                    register_b ^= program[ip + 1]
-                case 2:  # bst
-                    register_b = combo(ip + 1) & 0x7
-                case 3:  # jnz
-                    if register_a != 0:
-                        ip = program[ip + 1]
-                        continue
-                case 4:  # bxc
-                    register_b ^= register_c
-                case 5:  # out
-                    out.append(combo(ip + 1) & 7)
-                case 6:  # bdv
-                    register_b = register_a // 2 ** combo(ip + 1)
-                case 7:  # cdv
-                    register_c = register_a // 2 ** combo(ip + 1)
-            ip += 2
+        out = run_program(register_a, register_b, register_c, program)
 
         return ",".join(map(str, out))
 
     @classmethod
-    def part2(cls, input: str) -> str:
-        pass
+    def part2(cls, input: str) -> int:
+        numbers = re.findall(r"\d+", input)
+
+        _, register_b, register_c = map(int, numbers[:3])
+        program = list(map(int, numbers[3:]))
+
+        cur = [0]
+
+        # It came to me in a dream
+        for entry in reversed(program):
+            next_gen = []
+
+            for num in cur:
+                num *= 8
+                for n in range(8):
+                    output = run_program(num + n, register_b, register_c, program)
+                    result = output[0]
+                    if result == entry:
+                        next_gen.append(num + n)
+
+            cur = next_gen
+
+        return cur[0]

2024/tests/samples/17.txt renamed to 2024/tests/samples/17.1.txt

@@ -0,0 +1,5 @@
+Register A: 2024
+Register B: 0
+Register C: 0
+
+Program: 0,3,5,4,3,0

2024/tests/samples/17.2.txt

@@ -4,4 +4,8 @@
 
 
 def test_sample_part1() -> None:
-    assert DayRunner.part1(get_data(17)) == "4,6,3,5,6,3,5,2,1,0"
+    assert DayRunner.part1(get_data(17, 1)) == "4,6,3,5,6,3,5,2,1,0"
+
+
+def test_sample_part2() -> None:
+    assert DayRunner.part2(get_data(17, 2)) == 117440