Day 15

Author: monk-time

year2020/aoc15.py

@@ -1,15 +1,38 @@
 # https://adventofcode.com/2020/day/15
 
+from collections import defaultdict
+from collections.abc import Generator
+from itertools import islice
+
 from helpers import read_puzzle
 
+PART_1 = 2020
+PART_2 = 30000000
+
+
+def parse(s: str) -> tuple[int, ...]:
+    return tuple(int(x) for x in s.split(','))
+
+
+def gen_memory_game(start_nums: tuple[int, ...]) -> Generator[int]:
+    last_seen: defaultdict[int, int | None] = defaultdict(lambda: None)
+    turn, n = 0, 0
+    for turn, n in enumerate(start_nums, start=1):
+        last_seen[n] = turn
+        yield n
+    while True:
+        prev, last_seen[n] = last_seen[n], turn
+        yield (n := 0 if prev is None else turn - prev)
+        turn += 1
+
 
-def parse(s: str) -> list[int]:
-    return [int(line) for line in s.split()]
+def get_nth_num(start_nums: tuple[int, ...], n: int) -> int:
+    return next(islice(gen_memory_game(start_nums), n - 1, None))
 
 
 def solve() -> tuple[int, int]:
-    _puzzle = parse(read_puzzle())
-    return 0, 0
+    start_nums = parse(read_puzzle())
+    return get_nth_num(start_nums, PART_1), get_nth_num(start_nums, PART_2)
 
 
 if __name__ == '__main__':

year2020/tests/test_aoc15.py

@@ -1,5 +1,46 @@
-from year2020.aoc15 import solve
+from itertools import islice
+
+import pytest
+
+from year2020.aoc15 import PART_1, PART_2, gen_memory_game, get_nth_num, solve
+
+
+def test_gen_memory_game():
+    first_10 = tuple(islice(gen_memory_game((0, 3, 6)), 10))
+    assert first_10 == (0, 3, 6, 0, 3, 3, 1, 0, 4, 0)
+
+
+@pytest.mark.parametrize(
+    'start_nums, result',
+    (
+        ((0, 3, 6), 436),
+        ((1, 3, 2), 1),
+        ((2, 1, 3), 10),
+        ((1, 2, 3), 27),
+        ((2, 3, 1), 78),
+        ((3, 2, 1), 438),
+        ((3, 1, 2), 1836),
+    ),
+)
+def test_get_2020th_num(start_nums, result):
+    assert get_nth_num(start_nums, PART_1) == result
+
+
+@pytest.mark.parametrize(
+    'start_nums, result',
+    (
+        ((0, 3, 6), 175594),
+        ((1, 3, 2), 2578),
+        ((2, 1, 3), 3544142),
+        ((1, 2, 3), 261214),
+        ((2, 3, 1), 6895259),
+        ((3, 2, 1), 18),
+        ((3, 1, 2), 362),
+    ),
+)
+def test_get_30000000th_num(start_nums, result):
+    assert get_nth_num(start_nums, PART_2) == result
 
 
 def test_solve():
-    assert solve() == (0, 0)
+    assert solve() == (700, 51358)