add 2024/day22.2

2024/22/main.py

@@ -17,7 +17,11 @@ def mix_and_prune(i: int, j: int) -> int:
     return (i ^ j) % 16777216
 
 
-def calc_number(num: int) -> int:
+def calc_secret(num: int) -> (int, list, list):
+    changes = []
+    numbers = []
+    prev_num = num
+
     for _ in range(2000):
         # step 1
         tmp = num * 64
@@ -31,17 +35,57 @@ def calc_number(num: int) -> int:
         tmp = num * 2048
         num = mix_and_prune(num, tmp)
 
-    return num
+        numbers.append(num)
+        changes.append(num % 10 - prev_num % 10)
+        prev_num = num
+
+    return num, numbers, changes
+
+
+def find_in_list(lst: list, seq: list) -> int:
+    seq_len = 4
+    for i in range(len(lst) - seq_len + 1):
+        if lst[i : i + seq_len] == seq:
+            return i
+    return -1
+
+
+def get_subsequences(lsts):
+    sequences = set()
+    length = 4
+    for lst in lsts:
+        sequences.update(
+            tuple(lst[i : i + length]) for i in range(len(lst) - length + 1)
+        )
+    return sequences
 
 
 def main():
-    intial_numbers = readinput()
+    intial_secrets = readinput()
 
     # part 1
-    new_numbers = {}
-    for num in intial_numbers:
-        new_numbers[num] = calc_number(num)
-    print("Sum: %d" % sum(new_numbers.values()))
+    new_secrets = {}
+    for num in intial_secrets:
+        new_secrets[num], _, _ = calc_secret(num)
+    print("Sum: %d" % sum(new_secrets.values()))
+
+    # part 2 (very ugly and very slow brute force)
+    secret_list = {}
+    changes = {}
+    for num in intial_secrets:
+        _, secret_list[num], changes[num] = calc_secret(num)
+
+    bananas = {}
+    sequences = get_subsequences(changes.values())
+    for i, sequence in enumerate(sequences):
+        print(f"\rProcessing sequence {i}/{len(sequences) - 1}", end="", flush=True)
+        for secret, change_values in changes.items():
+            seq_pos = find_in_list(change_values, list(sequence))
+            if seq_pos != -1:
+                bananas[sequence] = (
+                    bananas.get(sequence, 0) + secret_list[secret][seq_pos + 3] % 10
+                )
+    print("\nBananas: %d " % max(bananas.values()))
 
 
 if __name__ == "__main__":