From d87b56e0ee069f6af50f14fb0f7d1fcdcd3e8d03 Mon Sep 17 00:00:00 2001
From: Stefan Harmuth <stefan.harmuth@net-m.de>
Date: Tue, 21 Dec 2021 11:11:08 +0100
Subject: [PATCH] day21

---
 day21.py                | 57 ++++++++++++++++++++---------------------
 inputs/test_input21_1_1 |  2 ++
 inputs/test_input21_2_1 |  2 ++
 3 files changed, 32 insertions(+), 29 deletions(-)
 create mode 100644 inputs/test_input21_1_1
 create mode 100644 inputs/test_input21_2_1

diff --git a/day21.py b/day21.py
index e0322d2..5075752 100644
--- a/day21.py
+++ b/day21.py
@@ -1,10 +1,32 @@
 from tools.aoc import AOCDay
+from tools.tools import cache
 from typing import Any
 
 
+# 1 way to roll a 3, 3 ways to roll a 4, ...
+ROLLS = [(1, 3), (3, 4), (6, 5), (7, 6), (6, 7), (3, 8), (1, 9)]
+
+
+@cache
+def roll(pos1: int, pos2: int, score1: int = 0, score2: int = 0) -> (int, int):
+    if score2 >= 21:
+        return 0, 1
+
+    p1wins = p2wins = 0
+
+    for sub_universes, pos_inc in ROLLS:
+        new_pos1 = 1 + (pos_inc + pos1 - 1) % 10
+        # notice the switch-a-roo!
+        sub_p2wins, sub_p1wins = roll(pos2, new_pos1, score2, score1 + new_pos1)
+        p1wins += sub_p1wins * sub_universes
+        p2wins += sub_p2wins * sub_universes
+
+    return p1wins, p2wins
+
+
 class Day(AOCDay):
-    test_solutions_p1 = [739785]
-    test_solutions_p2 = [444356092776315]
+    test_solutions_p1 = [739785, 551901]
+    test_solutions_p2 = [444356092776315, 272847859601291]
 
     def part1(self) -> Any:
         pos = list(map(int, [x.split(": ")[1] for x in self.getInput()]))
@@ -13,7 +35,7 @@ class Day(AOCDay):
         rolls = 0
         dice_roll = 0
         while scores[0] < target_score and scores[1] < target_score:
-            for player in [0, 1]:
+            for player in range(2):
                 this_roll = 0
                 rolls += 3
                 for _ in range(3):
@@ -27,30 +49,7 @@ class Day(AOCDay):
 
         return min(scores) * rolls
 
-    def roll_dirac(self, pos1, pos2, score1, score2, cache):
-        if (pos1, pos2, score1, score2) in cache:
-            return cache[(pos1, pos2, score1, score2)]
-        p1wins = 0
-        p2wins = 0
-        target_score = 21
-        for roll1 in range(3, 10):
-            tpos1 = 1 + (roll1 + pos1 - 1) % 10
-            if score1 + tpos1 >= target_score:
-                p1wins += 27
-            else:
-                for roll2 in range(3, 10):
-                    tpos2 = 1 + (roll2 + pos2 - 1) % 10
-                    if score2 + tpos2 >= target_score:
-                        p2wins += 27
-                    else:
-                        sub_p1wins, sub_p2wins = self.roll_dirac(tpos1, tpos2, score1 + tpos1, score2 + tpos2, cache)
-                        p1wins += sub_p1wins
-                        p2wins += sub_p2wins
-
-        cache[(pos1, pos2, score1, score2)] = p1wins, p2wins
-        return p1wins, p2wins
-
     def part2(self) -> Any:
-        pos = list(map(int, [x.split(": ")[1] for x in self.getInput()]))
-        wins1, wins2 = self.roll_dirac(pos[0], pos[1], 0, 0, {})
-        return wins1, wins2
+        pos1, pos2 = map(int, [x.split(": ")[1] for x in self.getInput()])
+        wins1, wins2 = roll(pos1, pos2)
+        return max(wins1, wins2)
diff --git a/inputs/test_input21_1_1 b/inputs/test_input21_1_1
new file mode 100644
index 0000000..5ad5c4a
--- /dev/null
+++ b/inputs/test_input21_1_1
@@ -0,0 +1,2 @@
+Player 1 starting position: 7
+Player 2 starting position: 3
diff --git a/inputs/test_input21_2_1 b/inputs/test_input21_2_1
new file mode 100644
index 0000000..5ad5c4a
--- /dev/null
+++ b/inputs/test_input21_2_1
@@ -0,0 +1,2 @@
+Player 1 starting position: 7
+Player 2 starting position: 3