Cleanup

08/8.py

@@ -0,0 +1,96 @@
+import numpy as np
+from pprint import pprint
+import sys
+
+from numpy.lib.function_base import delete
+
+EASY_LENGTHS = np.array([2, 4, 3, 7])
+
+vec_len = np.vectorize(len)
+
+
+def part_one(lines):
+
+    total = 0
+
+    for line in lines:
+        digits, output_str = line.split("|")
+        digits = digits.split()
+        output_value: np.ndarray = output_str.split()
+        lengths: np.ndarray = vec_len(output_value)
+        for length in lengths:
+            if length in EASY_LENGTHS:
+                total += 1
+
+    return total
+
+
+def part_two(lines):
+    total = 0
+
+    for line in lines:
+        digits, output_str = line.split("|")
+        digits = np.array([frozenset(list(x)) for x in digits.split()])
+
+        output_value: np.ndarray = np.array(
+            [frozenset(list(x)) for x in output_str.split()]
+        )
+
+        lengths: np.ndarray = vec_len(digits)
+
+        mapping = {}
+
+        # Get easy mappings
+        mapping[1] = digits[lengths == 2][0]
+        mapping[4] = digits[lengths == 4][0]
+        mapping[7] = digits[lengths == 3][0]
+        mapping[8] = digits[lengths == 7][0]
+
+        digits = digits[digits != mapping[1]]
+        digits = digits[digits != mapping[4]]
+        digits = digits[digits != mapping[7]]
+        digits = digits[digits != mapping[8]]
+        lengths: np.ndarray = vec_len(digits)
+
+        # 3: length 5, superset of 7
+        mapping[3] = digits[(lengths == 5) & (mapping[7] <= digits)][0]
+        digits = digits[digits != mapping[3]]
+        lengths: np.ndarray = vec_len(digits)
+
+        # 5: length 5, 3 parts of 4
+        digits_5 = digits[lengths == 5]
+        mapping[5] = digits_5[vec_len(digits_5 & mapping[4]) == 3][0]
+        digits = digits[digits != mapping[5]]
+        lengths: np.ndarray = vec_len(digits)
+
+        # 2: Remaining length 5
+        mapping[2] = digits[lengths == 5][0]
+        digits = digits[digits != mapping[2]]
+        del lengths  # make sure the invalid list isn't used anymore
+
+        # -- Only length 6 left
+
+        # 9: superset of 4
+        mapping[9] = digits[mapping[4] <= digits][0]
+        digits = digits[digits != mapping[9]]
+
+        # 0: superset of 7
+        mapping[0] = digits[mapping[7] <= digits][0]
+        digits = digits[digits != mapping[0]]
+
+        assert len(digits) == 1
+
+        # 6: remaining
+        mapping[6] = digits[0]
+
+        decode_mapping = {v: str(k) for k, v in mapping.items()}
+
+        total += int("".join([decode_mapping[x] for x in output_value]))
+
+    return total
+
+
+if __name__ == "__main__":
+    lines = sys.stdin.readlines()
+    print("Part one:", part_one(lines))
+    print("Part two:", part_two(lines))