day14.py

#!/usr/bin/env python3

import os
import sys
import random
import re
from collections import defaultdict
from itertools import permutations, product

src = open("input.txt", "r").readlines()

# src = """mask = XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X
# mem[8] = 11
# mem[7] = 101
# mem[8] = 0""".splitlines()

# src = """mask = 000000000000000000000000000000X1001X
# mem[42] = 100
# mask = 00000000000000000000000000000000X0XX
# mem[26] = 1""".splitlines()

src = [r.strip() for r in src]


def part1():
    mem = defaultdict(lambda: 0)
    xmask = 0
    y = 0

    for line in src:
        if line.startswith("mask"):
            mask = line.split(" = ")[1]
            xmask = int(mask.replace("1", "0").replace("X", "1"), 2)
            y = int(mask.replace("X", "0"), 2)
        else:
            m = re.search(r"mem\[(\d+)\] = (\d+)", line)
            addr = int(m.group(1))
            value = int(m.group(2))

            mem[addr] = (xmask & value) | y

    print(sum(mem.values()))


def part2():
    mem = defaultdict(lambda: 0)
    pvalues = []
    omask = 0

    for line in src:
        if line.startswith("mask"):
            mask = line.split(" = ")[1]

            omask = int(mask.replace("X", "0"), 2)

            bits = [1 << int(i) for i, x in enumerate(mask[::-1]) if x == "X"]

            perms = product([True, False], repeat=len(bits))

            pvalues = []
            for x in [p for p in perms]:
                pvalues.append(sum([bits[i] for i, b in enumerate(x) if x[i]]))
            pvalues.sort()

        else:
            m = re.search(r"mem\[(\d+)\] = (\d+)", line)
            addr = int(m.group(1))
            value = int(m.group(2))

            for pv in pvalues:
                mem[(addr | omask) ^ pv] = value

    print(sum(mem.values()))


part1()
part2()