Advent of Code 2020 Day 14

AoC2020_day14.py

import re

## Part 1

def day14_p1(program):
    mask = None
    mem = {}
    for line in program:
        lh, rh = line.split(" = ")
        if lh == "mask":
            mask = rh
        if lh.startswith("mem"):
            addr = re.match("^mem\[(\d+)\]", lh).groups()[0]
            argv = "{:036b}".format(int(rh))

            val = ""
            for m, v in zip(mask, argv):
                if m == "X":
                    val += v
                else:
                    val += m           
            mem[addr] = int(val, 2)

    return sum(mem.values())

with open("input_14") as program:
    print(day14_p1(program))

## Part 2

def get_addr(addr_mask):
    if "X" in addr_mask:
        for r in ("0", "1"):
            for addr in get_addr(addr_mask.replace("X", r, 1)):
                yield addr
    else:
        yield addr_mask

def day14_p2(program):
    mask = None
    mem = {}
    for line in program:
        lh, rh = line.split(" = ")
        if lh == "mask":
            mask = rh
        if lh.startswith("mem"):
            addr = re.match("^mem\[(\d+)\]", lh).groups()[0]
            addr = "{:036b}".format(int(addr))
            argv = int(rh)

            addr_m = ""
            for m, v in zip(mask, addr):
                if m == "0":
                    addr_m += v
                else:
                    addr_m += m

            for addr in get_addr(addr_m):
                mem[int(addr, 2)] = argv

    return sum(mem.values())

with open("input_14") as fp:
    program = [line for line in fp.readlines()]

    print("Part 1: ", day14_p1(program))
    print("Part 2: ", day14_p2(program))