Advent of Code 2024 Day 17

from pathlib import Path
from typing import Sequence
from operator import xor
from dataclasses import dataclass

TEST = """**REDACTED**"""


@dataclass
class Machine:
    reg_a: int
    reg_b: int
    reg_c: int
    program: list[int]

    def combo(self, operand: int) -> int:
        match operand:
            case 4:
                return self.reg_a
            case 5:
                return self.reg_b
            case 6:
                return self.reg_c
            case 7:
                raise RuntimeError("Reserved operand")
            case _:
                return operand

    def run(self, reg_a: int, reg_b: int = 0, reg_c: int = 0) -> list[int]:
        self.reg_a = reg_a
        self.reg_b = reg_b
        self.reg_c = reg_c
        return list(self.execute())

    def execute(self, pc: int = 0) -> Sequence[int]:
        program = self.program
        while pc < len(program):
            instr, operand = program[pc], program[pc + 1]
            pc += 2
            match instr:
                case 0:  # adv
                    self.reg_a = self.reg_a // (1 << self.combo(operand))
                case 1:  # bxl
                    self.reg_b = xor(self.reg_b, operand)
                case 2:
                    self.reg_b = self.combo(operand) % 8
                case 3:  # jnz
                    if self.reg_a != 0:
                        pc = operand
                case 4:
                    self.reg_b = xor(self.reg_b, self.reg_c)
                case 5:
                    yield self.combo(operand) % 8
                case 6:
                    self.reg_b = self.reg_a // (1 << self.combo(operand))
                case 7:
                    self.reg_c = self.reg_a // (1 << self.combo(operand))


def parse(input: str) -> Machine:
    lines = input.splitlines()
    reg_a = int(lines[0].split(":")[1].strip())
    reg_b = int(lines[1].split(":")[1].strip())
    reg_c = int(lines[2].split(":")[1].strip())
    program = [int(n) for n in lines[4].split(":")[1].strip().split(",")]
    return Machine(reg_a, reg_b, reg_c, program)


def part1(input: str) -> list[int]:
    device = parse(input)
    print(device)
    return list(device.execute())


assert part1(TEST) == [4, 6, 3, 5, 6, 3, 5, 2, 1, 0]

INPUT = Path("input/day17.txt").read_text()
result1 = part1(INPUT)
print("part 1:", ",".join(str(r) for r in result1))


def solve(device: Machine, reg_a: int, required: list[int]) -> Sequence[int]:
    reg_a *= 8
    for n in range(8):
        output = device.run(reg_a + n)
        if output == required:
            yield reg_a + n


def part2(input: str) -> list[int]:
    device = parse(input)
    required = []

    reg_a_options = {0}
    for value in reversed(device.program):
        required.insert(0, value)

        reg_a_options = {
            a for reg_a in reg_a_options for a in solve(device, reg_a, required)
        }

    for a in reg_a_options:
        print(a, device.run(a))
    return min(reg_a_options)


part2_total = part2(INPUT)
print(f"{part2_total=:,}")  # 37,221,274,271,220