def read_file(): with open("input.txt","r") as f: return f.read().split('\n') def process(instructions): mask = "X"*36 memory = {} for instruction in instructions: if instruction[0:4] == "mask": mask = instruction[7:] # trim "mask = " else: address = instruction[4:].split(']')[0] # trim "mem[", then split off everyting after "]" value = instruction.split('=')[1][1:] # split, take the second half, and trim the space value = str(format(int(value),'036b')) # 36 bits of binary, with leading 0s result = "" for m, b in zip(mask, value): if m == "0" or m == "1": # mask is applied result += m # so rewrite with mask bit else: # m == "X" result += b # so keep value bit memory[address] = result return memory def part1(instructions): memory = process(instructions) result = 0 for address in memory: result += int(memory[address],2) return result def addresses(address, mask): addresses = [] # TODO: learn how to use itertools to avoid having to manually do this stuff? bc "there must be a better way" ~raymond hettinger temp = "" for i in range(len(mask)): if mask[i] == "0": temp += address[i] elif mask[i] == "1": temp += mask[i] else: # mask[i] == "X" temp += "X" addresses.append(temp) iterations = mask.count("X") for i in range(iterations): # replace each X one by one a0s = [] a1s = [] for a in addresses: a0 = a.replace("X","0",1) # replace first X with 0 a1 = a.replace("X","1",1) # replace first X with 1 a0s.append(a0) a1s.append(a1) addresses = a0s + a1s # repopulate list -- should double in size return addresses def process2(instructions): mask = "0"*36 memory = {} for instruction in instructions: if instruction[0:4] == "mask": mask = instruction[7:] else: address = instruction[4:].split(']')[0] address = str(format(int(address),'036b')) value = instruction.split('=')[1][1:] value = str(format(int(value),'036b')) exploded_addrs = addresses(address, mask) # explode mask due to quantum bits for address in exploded_addrs: memory[address] = value return memory def part2(instructions): memory = process2(instructions) result = 0 for address in memory: result += int(memory[address],2) return result def main(): instructions = read_file() print(part1(instructions), part2(instructions)) if __name__ == "__main__": main()