From b23676bf04cd24b74df5955e47bd32a38c72bae2 Mon Sep 17 00:00:00 2001
From: Bert Peters <bert@bertptrs.nl>
Date: Thu, 26 Dec 2024 16:43:15 +0100
Subject: [PATCH] Implement 2024 day 21 part 2

Invent a unit test for it too, because why not, it's Christmas
---
 2024/src/aoc/days/day21.py | 31 +++++++++++++++++++++++++++++--
 2024/tests/test_day21.py   |  4 ++++
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/2024/src/aoc/days/day21.py b/2024/src/aoc/days/day21.py
index 4f1f5b6..6d3d6c1 100644
--- a/2024/src/aoc/days/day21.py
+++ b/2024/src/aoc/days/day21.py
@@ -1,4 +1,5 @@
 import functools
+from collections import Counter, defaultdict
 
 from . import SeparateRunner
 
@@ -114,6 +115,18 @@ def decode(code: str, pad: dict[str, tuple[int, int]]) -> str:
     return result
 
 
+def count_steps(path: str, count: int) -> dict[str, int]:
+    cur = "A"
+    counts = defaultdict(int)
+
+    for c in path:
+        step = shortest_dirpad(cur, c)
+        cur = c
+        counts[step] += count
+
+    return counts
+
+
 class DayRunner(SeparateRunner):
     @classmethod
     def part1(cls, input: str) -> int:
@@ -128,5 +141,19 @@ class DayRunner(SeparateRunner):
         return result
 
     @classmethod
-    def part2(cls, input: str) -> int:
-        pass
+    def part2(cls, input: str, robots=25) -> int:
+        result = 0
+        for code in input.strip().split("\n"):
+            numpad = encode_shortest_numpad(code)
+            keypresses = Counter([numpad])
+
+            for _ in range(robots + 1):
+                new_presses = Counter()
+                for subroute, count in keypresses.items():
+                    new_presses.update(count_steps(subroute, count))
+
+                keypresses = new_presses
+
+            result += int(code[:-1]) * keypresses.total()
+
+        return result
diff --git a/2024/tests/test_day21.py b/2024/tests/test_day21.py
index eaddd02..9dda72a 100644
--- a/2024/tests/test_day21.py
+++ b/2024/tests/test_day21.py
@@ -49,3 +49,7 @@ def test_encode_shortest_dirpad_twice(code: str, answer: str) -> None:
 
 def test_sample_part1() -> None:
     assert DayRunner.part1(get_data(21)) == 126384
+
+
+def test_sample_part2() -> None:
+    assert DayRunner.part2(get_data(21), robots=2) == 126384