Advent of Code Day 14

"""Day 14 of Advent of Code 2020 Solution"""
from collections.abc import MutableMapping


class BinArray(MutableMapping):
    def __init__(self, mask=None, mem_len=36):
        if mask:
            self._mask = mask
        self.mem_len = mem_len
        self.mem = dict()

    @property
    def mask(self):
        return self._mask

    @mask.setter
    def mask(self, mask):
        _mask = {}
        for i, v in enumerate(reversed(mask)):
            if v != "X":
                _mask[i] = v
        self._mask = _mask

    def __getitem__(self, item):
        return self.mem[item]

    def __setitem__(self, key, value):
        if self.mask:
            value = bin(value)[2:].zfill(self.mem_len)
            new_value = []
            for i, v in enumerate(reversed(value)):
                if i in self.mask:
                    new_value.append(self.mask[i])
                else:
                    new_value.append(v)
            self.mem[key] = "".join([*reversed(new_value)])
        else:
            self.mem[key] = value

    def __delitem__(self, key):
        del self.mem[key]

    def __iter__(self):
        for key in self.mem.keys():
            yield key

    def dec_values(self):
        for k in self:
            yield int(self[k], 2)

    def __len__(self):
        return len(self.mem)


class DecoderChip:
    def __init__(self):
        self.mem = dict()
        self._mask = None
        self.mem_len = 36

    @property
    def mask(self):
        return self._mask

    @mask.setter
    def mask(self, mask):
        self._mask = "".join([*reversed(mask)])

    def __setitem__(self, key, value):
        mem_locations = [""]
        key = bin(key)[2:].zfill(self.mem_len)
        for i, v in enumerate(reversed(key)):
            if self.mask[i] == "X":
                mem_floats_zero = mem_locations.copy()
                for j, mem_float in enumerate(mem_floats_zero):
                    mem_floats_zero[j] += "0"
                mem_floats_one = mem_locations.copy()
                for j, mem_float in enumerate(mem_floats_one):
                    mem_floats_one[j] += "1"
                mem_locations = mem_floats_zero + mem_floats_one
            elif self.mask[i] == "1":
                for j, mem_float in enumerate(mem_locations):
                    mem_locations[j] += "1"
            else:
                for j, location in enumerate(mem_locations):
                    mem_locations[j] += v
        mem_locations = ["".join([*reversed(location)]) for location in mem_locations]
        for location in mem_locations:
            self.mem[location] = value


def part_a(file_location):
    mem = BinArray()
    with open(file_location, "r") as f:
        for line in f.readlines():
            line = line.split(" = ")
            if line[0] == "mask":
                mem.mask = line[1].strip()
            else:
                mem_slot, value = line[0][4:-1], int(line[1])
                mem[mem_slot] = value
    return sum([*mem.dec_values()])


def part_b(file_location):
    decoder = DecoderChip()
    with open(file_location, "r") as f:
        for line in f.readlines():
            line = line.split(" = ")
            if line[0] == "mask":
                decoder.mask = line[1].strip()
            else:
                mem_slot, value = int(line[0][4:-1]), int(line[1])
                decoder[mem_slot] = value
    return sum([*decoder.mem.values()])


if __name__ == "__main__":
    file_location = r"data/day14.txt"
    print(part_a(file_location))
    print(part_b(file_location))