2024: d22: ex2: add solution

2024/d22/ex2/ex2.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python
+
+import itertools
+import sys
+from collections import deque
+from collections.abc import Iterator
+
+Sequence = tuple[int, int, int, int]
+
+
+def solve(input: str) -> int:
+    def monkey_hash(seed: int) -> int:
+        MASK = (1 << 24) - 1
+        seed ^= seed << 6
+        seed &= MASK
+
+        seed ^= seed >> 5
+        seed &= MASK
+
+        seed ^= seed << 11
+        seed &= MASK
+
+        return seed
+
+    def list_prices(seed: int) -> list[int]:
+        prices: list[int] = [seed % 10]
+        for _ in range(2000):
+            seed = monkey_hash(seed)
+            prices.append(seed % 10)
+        return prices
+
+    def find_sequences(prices: list[int]) -> dict[Sequence, int]:
+        # Adapted from an `itertools` recipe
+        def sliding_window() -> Iterator[tuple[int, ...]]:
+            iterator = iter(prices)
+            window = deque(itertools.islice(iterator, 5 - 1), maxlen=5)
+            for x in iterator:
+                window.append(x)
+                yield tuple(window)
+
+        assert len(prices) == 2001  # Sanity check
+        sequences: dict[Sequence, int] = {}
+        for i, (a, b, c, d, e) in enumerate(sliding_window(), start=4):
+            sequences.setdefault((b - a, c - b, d - c, e - d), i)
+        return sequences
+
+    def run_sequences(
+        seq: Sequence, prices: list[list[int]], sequences: list[dict[Sequence, int]]
+    ) -> int:
+        return sum(
+            price[sequence[seq]]
+            for price, sequence in zip(prices, sequences)
+            if seq in sequence
+        )
+
+    def all_sequences(sequences: list[dict[Sequence, int]]) -> set[Sequence]:
+        res: set[Sequence] = set()
+        for seq in sequences:
+            res |= seq.keys()
+        return res
+
+    seeds = [int(n) for n in input.splitlines()]
+    prices = [list_prices(seed) for seed in seeds]
+    sequences = [find_sequences(price_list) for price_list in prices]
+    return max(
+        run_sequences(seq, prices, sequences) for seq in all_sequences(sequences)
+    )
+
+
+def main() -> None:
+    input = sys.stdin.read()
+    print(solve(input))
+
+
+if __name__ == "__main__":
+    main()