Day 17, part 2

day17.py

import itertools


program = [2, 4, 1, 3, 7, 5, 0, 3, 1, 5, 4, 1, 5, 5, 3, 0]
rp = list(reversed(program))


def calc(ta):
    b = (ta & 0b111) ^ 0b011  # Look at the 3 rightmost bits of "ta"
    c = (
        ta >> b
    ) & 0b111  # Shift-right at most 7 bits, and look at the rightmost 3 bits.
    return b ^ 5 ^ c & 0b111


def create_large_num(l):
    n = l[0]
    for x in l[1:]:
        n = (n << 3) | x

    return n


def try_candidate(a):
    output = []

    while a > 0:
        ta = a & 0b1111111111  # Only use rightmost 10 bits
        a = a >> 3  # Shift right "a" by 3 bits (divide by 8).
        output.append(memo[ta])
        if program[: len(output)] != output:
            break  # Mismatch detected.

    return output


memo = {}
for n in range(1024):
    memo[n] = calc(n)

# Invert memo.
lookup = {}
for p in program:
    if not (p in lookup):
        lookup[p] = [k for k, v in memo.items() if v == p]

candidates = [[] for _ in program]

# Leftmost digit.
candidates[0] = list(filter(lambda x: x < 8, lookup[rp[0]]))

for i in range(1, len(rp)):
    # The next 3 bits can be found by first shifting the last number left by 3,
    # and then selecting from all potential candidates:
    # - Clear the lowest 3 bits of the next number (bitmask 0x3f8).
    # - Compare this number with the previous number shifted left by 3 (use the
    #   bitmask 0x3ff to only compare the lowest 10 bits). Note that low 3 bits
    #   will be zero, due to the shift.
    # - Use the low 3 bits of any such found number to create the list of
    #   candidates for this output position.
    for prev in candidates[i - 1]:
        prev = prev << 3
        for n in filter(lambda x: x & 0x3f8 == prev & 0x3ff, lookup[rp[i]]):
            candidates[i].append(prev | n & 0b111)

# Remove duplicates
for i in range(0, len(program) - 1):
    candidates[i] = list(set(candidates[i]))

# Sort candidates
for c in candidates:
    c.sort()

# The answer will actually be found in the first iteration of this loop.
for count, e in enumerate(itertools.product(*candidates)):
    n = create_large_num(list(e))
    output = try_candidate(n)
    if output == program:
        print("Answer:", n)
        break