Update day14.py

2020/Python/day14.py

@@ -1,3 +1,6 @@
+import itertools
+
+
 def parse_data() -> list[str]:
     lines = []
     with open("input.txt") as file:
@@ -7,7 +10,7 @@ def parse_data() -> list[str]:
 
 
 def solve_p1(program: list[str]) -> int:
-    bit_mask = "XXXXXXXXXXXXXXXXXXX"
+    bit_mask = ""
     mem = {}
 
     def apply_mask(_number: int, _mask: str) -> int:
@@ -30,4 +33,41 @@ def solve_p1(program: list[str]) -> int:
     return sum(mem.values())
 
 
+def solve_p2(program: list[str]) -> int:
+    bit_mask = ""
+    mem = {}
+
+    def apply_mask(_address: str, _mask: str) -> list[str]:
+        res = ['0' if x not in "01" else x for x in format(int(_address), "36b")]
+        for i in range(len(_mask)):
+            if _mask[i] == "1":
+                res[i] = "1"
+            elif _mask[i] == "X":
+                res[i] = "X"
+        return res
+
+    def resolve_floating(_address: list[str]) -> list[int]:
+        combinations = list(itertools.product(["0", "1"], repeat=_address.count("X")))
+        res = []
+        for combo in combinations:
+            res.append("".join(_address))
+            for c in combo:
+                res[-1] = res[-1].replace('X', c, 1)
+        return res
+
+    for line in program:
+        lvalue, rvalue = line.split(" = ")
+        if lvalue == "mask":
+            bit_mask = rvalue
+        elif lvalue[:3] == "mem":
+            address = lvalue[4:-1]
+            result = apply_mask(address, bit_mask)
+            floating_addresses = resolve_floating(result)
+            for floating_address in floating_addresses:
+                mem[floating_address] = int(rvalue)
+
+    return sum(mem.values())
+
+
 print(solve_p1(parse_data()))
+print(solve_p2(parse_data()))