aoc_2024_day_17.py

text = """Register A: 6819585
Register B: 0
Register C: 0

Program: 2,4,1,5,7,5,1,6,4,1,5,5,0,3,3,0"""


# can't assign inside lambdas
def regset(regs, key, val):
    regs[key] = val

ops = [
    ("adv", True,  lambda regs, x: regset(regs, "A", regs["A"] >> x)),
    ("bxl", False, lambda regs, x: regset(regs, "B", regs["B"] ^ x)),
    ("bst", True,  lambda regs, x: regset(regs, "B", x & 7)),
    ("jnz", False, lambda regs, x: regset(regs, "IP", x if regs["A"] else regs["IP"])),
    ("bxc", False, lambda regs, x: regset(regs, "B", regs["B"] ^ regs["C"])),
    ("out", True,  lambda regs, x: regset(regs, "OUT", regs["OUT"] + [x & 7])),
    ("bdv", True,  lambda regs, x: regset(regs, "B", regs["A"] >> x)),
    ("cdv", True,  lambda regs, x: regset(regs, "C", regs["A"] >> x)),
]


### PARSE INPUT

regs = dict()
prog = list()

regs_input, program_input = text.split("\n\n")

for line, reg in zip(regs_input.splitlines(), "ABC"):
    assert(line.startswith(f"Register {reg}: "))
    regs[reg] = int(line.split(": ")[1])

assert(program_input.startswith(f"Program: "))
prog = [int(x) for x in program_input[len(f"Program: "):].split(",")]

print(regs)
print(prog)


### PART 1

def run(prog, regs):
    regs["IP"] = 0
    regs["OUT"] = []

    while regs["IP"] < len(prog):
        ip = regs["IP"]
        opcode = prog[ip]
        operand = prog[ip+1]

        _, use_combo, fn = ops[opcode]

        if use_combo and operand >= 4:
            operand = regs["ABC"[operand - 4]]
        
        fn(regs, operand)
        
        if regs["IP"] == ip: # bump ip ptr iff didn't jump
            regs["IP"] += 2

    return regs["OUT"]

print(run(prog, regs))


## PART2

def get_bit(num, bit):
    return (num >> bit) & 1
    
# generate possible 3-digit combinations at a given offset
def schroedingdong(digits, offset):
    def get_possible_digits(digits, offset):
        if offset in digits:
            return [digits[offset]]
        else:
            return [0, 1]
    combinations = [()]
    for i in range(3):
        new_combos = []
        for k in get_possible_digits(digits, offset+i):
            new_combos += [(k,) + c for c in combinations]
        combinations = new_combos
    return combinations

# the last 3 digits added might or might not work
def check_last_digits_validity(digits, offset, tgt):
    b0 = digits[offset + 0]
    b1 = digits[offset + 1]
    b2 = digits[offset + 2]

    b = b2 * 4 + b1 * 2 + b0
    c = b ^ 3 ^ tgt
    off = (b ^ 5) + offset
    
    c0 = get_bit(c, 0)
    c1 = get_bit(c, 1)
    c2 = get_bit(c, 2)

    if digits.get(off + 0, c0) != c0: return False
    if digits.get(off + 1, c1) != c1: return False
    if digits.get(off + 2, c2) != c2: return False

    digits[off + 0] = c0
    digits[off + 1] = c1
    digits[off + 2] = c2

    return True

def generate_values_for_a(digits, offset, target_output):
    if not len(target_output):
        # put digits together
        num = 0
        for k, v in digits.items():
            num |= v << k
        yield num

    else:
        # try all (still) possible combinations of 3 digits:
        for b2, b1, b0 in schroedingdong(digits, offset):
            digits2 = dict(digits)

            digits2[offset + 0] = b0
            digits2[offset + 1] = b1
            digits2[offset + 2] = b2

            if check_last_digits_validity(digits2, offset, target_output[0]):
                for result in generate_values_for_a(digits2, offset + 3, target_output[1:]):
                    yield result

possible_values_for_a = [x for x in generate_values_for_a(digits=dict(), offset=0, target_output=prog)]
print(possible_values_for_a)

smallest_possible_a = min(possible_values_for_a)
print(f"A = {smallest_possible_a}")

# check result
regs["A"] = smallest_possible_a
print(prog)
print(run(prog, regs))