2021 Day 21 Part 2

2021/day21/main.go

@@ -10,6 +10,7 @@ func main() {
 	lines := aoc.Input().ReadLines("day21.txt")
 	p1, p2 := parseInput(lines)
 	fmt.Println(part1(p1, p2))
+	fmt.Println(part2(p1, p2))
 }
 
 func parseInput(lines []string) (p1, p2 int) {
@@ -88,3 +89,116 @@ func (d *DD100) Roll() int {
 	d.val = (d.val + 1) % 100
 	return val + 1
 }
+
+func part2(p1, p2 int) int {
+	active := make(StateCounts)
+	active[State{p1, p2, 0, 0}] = 1
+
+	done := make(map[State]int)
+
+	for len(active) > 0 {
+		state, count := active.Pop()
+		nextActive, nextDone := state.Successors()
+		for next, extra := range nextActive {
+			active[next] += count * extra
+		}
+		for next, extra := range nextDone {
+			done[next] += count * extra
+		}
+	}
+
+	w1, w2 := 0, 0
+	for state, count := range done {
+		if state.Win1() {
+			w1 += count
+		} else {
+			w2 += count
+		}
+	}
+	return max(w1, w2)
+}
+
+type StateCounts map[State]int
+
+func (sc StateCounts) Pop() (State, int) {
+	for state, count := range sc {
+		delete(sc, state)
+		return state, count
+	}
+	panic("Pop from empty map")
+}
+
+func (sc StateCounts) Total() (count int) {
+	for _, c := range sc {
+		count += c
+	}
+	return
+}
+
+func max(a, b int) int {
+	if a > b {
+		return a
+	}
+	return b
+}
+
+type State struct {
+	p1, p2 int
+	s1, s2 int
+}
+
+func (s State) Done() bool {
+	return s.s1 >= 21 || s.s2 >= 21
+}
+
+func (s State) Win1() bool {
+	return s.s1 >= 21
+}
+
+var diracDie = []int{1, 2, 3}
+var diracRolls = makeDiracRolls()
+
+func makeDiracRolls() (rolls []int) {
+	for _, r1 := range diracDie {
+		for _, r2 := range diracDie {
+			for _, r3 := range diracDie {
+				rolls = append(rolls, r1+r2+r3)
+			}
+		}
+	}
+	return
+}
+
+func (s State) Successors() (active StateCounts, done StateCounts) {
+	active, done = make(StateCounts), make(StateCounts)
+	for _, r1 := range diracRolls {
+		t1 := s.P1(r1)
+		if t1.Done() {
+			done[t1]++
+			continue
+		}
+		for _, r2 := range diracRolls {
+			t2 := t1.P2(r2)
+			if t2.Done() {
+				done[t2]++
+				continue
+			}
+			active[t2]++
+		}
+	}
+	return
+}
+
+func (s State) P1(roll int) State {
+	next := s
+	next.p1 = move(s.p1, roll)
+	next.s1 = s.s1 + score(next.p1)
+	return next
+}
+
+func (s State) P2(roll int) State {
+	next := s
+	next.p2 = move(s.p2, roll)
+	next.s2 = s.s2 + score(next.p2)
+	return next
+}