Advent of code 2020, Day 14, Python3.9

day14.py

from sys import stdin
from functools import reduce, partial

def parse_mask(l):
    m = l.split('=')[1].strip()
    r = lambda v: (
        lambda s,i: 2*s+(v==i))
    return (reduce(r('1'),m,0),
            reduce(r('X'),m,0))

def parse_addrval(l):
    lh,rh = l.split('=')
    lh = lh.split(
        '[')[1].split(']')[0]
    return int(lh),int(rh)

def fl(mask):
    if not mask:
        return [0]
    lsb = mask & -mask
    z = fl(mask^lsb)
    return z+[m|lsb for m in z]

def apply(part1,st,l):
    base,mask,stor = st
    if l[3] == 'k':
        return (*parse_mask(l), stor)
    addr,val = parse_addrval(l)
    if part1:
        stor |= {addr: base|(val&mask)}
    else:
        stor |= {base|(addr^m): val
                 for m in fl(mask)}
    return (base,mask,stor)

def solve(I, part1):
    return (
        sum(reduce(
            partial(apply,
                    part1),
            I,(0,0,{}))[2].values()))


I = [l for l in stdin]
print (solve(I,True), solve(I, False))