AOC2020Day14inPython

data = [l.split(" = ") for l in open("data14.txt").readlines()]

# For Part 2 - this generates all the possible addresses given base and list of floating bits
def setbit(n, b):   return n | 1<<b
def clearbit(n, b): return n & ~(1<<b)
def allfloats(base, fl):
    if fl:
        yield from allfloats(clearbit(base,fl[0]), fl[1:])
        yield from allfloats(  setbit(base,fl[0]), fl[1:])
    else:
        yield base

mem1, mem2 = {}, {}
for (op, val) in data:
    if op == "mask":
        # mask = 000X0001010X01111X110X1X011X01110100
        setmask   = int(val.replace('X','0'),2)
        clearmask = int(val.replace('X','1'),2)
        fbits = [35-i for (i,c) in enumerate(val) if c=='X']
    else:
        # mem[6448] = 416920615
        loc, val = int(op[4:-1]), int(val)
        mem1[loc] = (val | setmask) & clearmask
        for l in allfloats(loc | setmask, fbits): mem2[l] = val

print("Part 1: ",sum(mem1.values()))
print("Part 2: ",sum(mem2.values()))