day14.py

def mask_val(mask, val):
    val = format(val, "#038b")[2:]
    new_val = ""
    for pair in zip(val, mask):
        if pair[1] != "X":
            new_val += pair[1]
        else:
            new_val += pair[0]
    return int(new_val, 2)

mask = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X"
assert mask_val(mask, 11) == 73
assert mask_val(mask, 101) == 101
assert mask_val(mask, 0) == 64

import re
with open("test_input.txt") as file:
    test_input = file.read().splitlines()

def run_it(data):
    mem = {}
    for cmd in data:
        if "mask" in cmd:
            mask = cmd.split(" = ")[1]
        else:
            idx, val = re.fullmatch(r"mem\[(\d+)\] = (\d+)", cmd).groups()
            mem[int(idx)] = mask_val(mask, int(val))
    return sum(mem.values())

assert run_it(test_input) == 165

with open("part1.txt") as file:
    data = file.read().splitlines()

run_it(data)

with open("part2_test.txt") as file:
    part2_test = file.read().splitlines()

def mask_val2(mask, val):
    val = format(val, "#038b")[2:]
    new_val = ""
    for pair in zip(val, mask):
        bit, mask_bit = pair
        if mask_bit == "0":
            new_val += str(bit)
        elif mask_bit == "1":
            new_val += "1"
        else:
            new_val += "X"
    return new_val

def expand_floaters(val, vals=[]):
    if "X" in val:
        expand_floaters(val.replace('X', '0', 1), vals)
        expand_floaters(val.replace('X', '1', 1), vals)
    else:
        vals.append(val)
    return vals

def part2(data):
    mem = {}
    for cmd in data:
        if "mask" in cmd:
            mask = cmd.split(" = ")[1]
        else:
            idx, val = re.fullmatch(r"mem\[(\d+)\] = (\d+)", cmd).groups()
            for mem_idx in expand_floaters(mask_val2(mask, int(idx)), []):
                mem[int(mem_idx, 2)] = int(val)
    return sum(mem.values())
    # return mem

assert part2(part2_test) == 208

part2(data)