aoc-2020/14.py

lines = []

with open("14.input") as f:
    for line in f.readlines():
        lines.append(line.strip())

andmask = ormask = 0
mask = ""
mem = {}

for line in lines:
    if line[0:4] == "mask":
        mask = line.split(" ")[2]

        andmask = int(mask.replace("X", "1"), 2)
        ormask = int(mask.replace("X", "0"), 2)
    else:
        line = line.split(" ")
        addr = int(line[0][4:-1])
        val_orig = int(line[2])
        val = (val_orig & andmask) | ormask
        mem[addr] = val

print("Answer 1: {}".format(sum(mem.values())))

mem = {}
floating = []

for line in lines:
    if line[0:4] == "mask":
        mask = line.split(" ")[2]
        ormask = int(mask.replace("X", "0"), 2)

        floating = []
        for (i, c) in enumerate(mask):
            if c == 'X':
                floating.append(35 - i)
    else:
        line = line.split(" ")
        addr = int(line[0][4:-1])
        addr |= ormask
        val = int(line[2])
        for num in range(0, 2 ** len(floating)):
            my_addr = addr
            for (j, index) in enumerate(floating):
                if (num >> j) & 1:
                    my_addr |= 1 << index
                else:
                    my_addr &= (0xFFFFFFFFF ^ (1 << index))

            mem[my_addr] = val

print("Answer 2: {}".format(sum(mem.values())))