aoc2024day24.py

# === Parsing ===

values = {}
dependencies = {}
zs = set()

with open("input.txt", "r") as f:
    for line in iter(f.readline, "\n"):
        wire, value = line.strip().split(": ")
        values[wire] = int(value)

    for line in f:
        inp1, op, inp2, _, out = line.split()
        dependencies[out] = (inp1, op, inp2)
        if out.startswith("z"):
            zs.add(out)


# === Part 1 ===
import operator

OPS = {
    "AND": operator.and_,
    "OR": operator.or_,
    "XOR": operator.xor,
}


def propagate(wire, dependencies, values):
    if wire not in values:
        inp1, op, inp2 = dependencies[wire]
        v1 = propagate(inp1, dependencies, values)
        v2 = propagate(inp2, dependencies, values)
        values[wire] = OPS[op](v1, v2)

    return values[wire]


def z_values_to_int(z_values):
    return int(
        "".join(str(z_values[z]) for z in sorted(z_values.keys(), reverse=True)),
        2,
    )


z_values = {z: propagate(z, dependencies, values) for z in zs}

print(z_values_to_int(z_values))


# === Part 2 ===
from itertools import chain, combinations
import random


def test_addition(bits, dependencies):
    xv = random.randint(0, pow(2, bits) - 1)
    bin_x = bin(xv).removeprefix("0b").zfill(bits)
    values = {f"x{idx:02}": int(bit) for idx, bit in enumerate(reversed(bin_x))}
    yv = random.randint(0, pow(2, bits) - 1)
    bin_y = bin(yv).removeprefix("0b").zfill(bits)
    values |= {f"y{idx:02}": int(bit) for idx, bit in enumerate(reversed(bin_y))}
    expected = (xv + yv) % pow(2, bits)

    zs = {f"z{idx:02}" for idx in range(bits)}
    z_values = {z: propagate(z, dependencies, values) for z in zs}
    z = z_values_to_int(z_values)
    return z == expected


def test_bits(bits, dependencies, n_trials):
    return all(test_addition(bits, dependencies) for _ in range(n_trials))


def get_all_dependencies(wire, dependencies):
    all_deps = set()
    to_explore = [(wire, set())]
    while to_explore:
        w, came_before = to_explore.pop()
        all_deps.add(w)
        if w in dependencies:
            inp1, _, inp2 = dependencies[w]
            if inp1 in came_before or inp2 in came_before:
                return False, set()
            to_explore.append((inp1, came_before | {w}))
            to_explore.append((inp2, came_before | {w}))

    return True, all_deps


def swap_dependencies(wires_to_swap, dependencies):
    w1, w2 = wires_to_swap
    temp = dependencies[w1]
    dependencies[w1] = dependencies[w2]
    dependencies[w2] = temp


MAX_BITS = 45

all_wires = set(dependencies.keys())
correct_wires = {wire for wire in values if wire[0] in "xy"}

print(get_all_dependencies("z03", dependencies))

swaps = []
for bits in range(1, MAX_BITS + 1):
    print(f"Testing {bits} ", end="")
    is_working = test_bits(bits, dependencies, 200)
    if is_working:
        print("fine!")
        _, all_deps = get_all_dependencies(f"z{bits - 1:02}", dependencies)
        correct_wires |= all_deps
        continue

    # Try random swaps until it works.
    incorrect_wires = all_wires - correct_wires
    high_inputs = {f"x{idx:02}" for idx in range(bits, MAX_BITS + 1)}
    high_inputs |= {f"y{idx:02}" for idx in range(bits, MAX_BITS + 1)}
    for wires_to_swap in combinations(incorrect_wires, 2):
        swap_dependencies(wires_to_swap, dependencies)
        # Am I depending on a high input?
        valid_deps = True
        for nbit in range(bits):
            is_valid, all_deps = get_all_dependencies(f"z{nbit:02}", dependencies)
            if not is_valid or all_deps & high_inputs:
                valid_deps = False
                break
        if not valid_deps:
            swap_dependencies(wires_to_swap, dependencies)
            continue

        is_working = test_bits(bits, dependencies, 200)
        if is_working:
            print(wires_to_swap, "worked!")
            swaps.append(wires_to_swap)
            break

        # If it didn't work, swap it back.
        swap_dependencies(wires_to_swap, dependencies)


print(",".join(sorted(chain.from_iterable(swaps))))