AOC2021 Day 21 (Python)

day21_tests.py

def play(player_a_start, player_b_start):
    a_score = 0
    b_score = 0
    a_pos = player_a_start
    b_pos = player_b_start

    die_rolls = 0
    next_die_value = 0

    def dice():
        nonlocal next_die_value
        nonlocal die_rolls

        result = next_die_value + 1
        die_rolls += 1
        next_die_value = result % 100
        return result

    while True:
        d1, d2, d3 = (dice(), dice(), dice())
        a_pos = ((a_pos + d1 + d2 + d3 - 1) % 10) + 1
        a_score = a_score + a_pos
        if a_score >= 1000:
            return die_rolls * b_score

        d1, d2, d3 = (dice(), dice(), dice())
        b_pos = ((b_pos + d1 +d2 + d3 - 1) % 10) + 1
        b_score = b_score + b_pos
        if b_score >= 1000:
            return die_rolls * a_score


def test_example():
    assert play(4, 8) == 739785


def test_deterministic():
    print()
    print(play(7, 3))


def group_and_count(l):
    counts = dict()
    for i in l:
        counts[i] = counts.get(i, 0) + 1
    return counts


dice_distribution = group_and_count(a + b + c for a in range(1, 4) for b in range(1, 4) for c in range(1, 4))


class GameState(object):

    def __init__(self, turn_count, player_a_pos, player_a_score, player_b_pos, player_b_score):
        self.turn_count = turn_count
        self.player_a_pos = player_a_pos
        self.player_a_score = player_a_score
        self.player_b_pos = player_b_pos
        self.player_b_score = player_b_score

    def next_states(self):
        result = dict()
        new_turn_count = self.turn_count + 1
        for dice_total, occurrence_count in dice_distribution.items():
            if self.turn_count % 2 == 0:
                new_player_a_pos = ((self.player_a_pos + dice_total - 1) % 10) + 1
                new_player_a_score = self.player_a_score + new_player_a_pos
                new_player_b_pos = self.player_b_pos
                new_player_b_score = self.player_b_score
            else:
                new_player_a_pos = self.player_a_pos
                new_player_a_score = self.player_a_score
                new_player_b_pos = ((self.player_b_pos + dice_total - 1) % 10) + 1
                new_player_b_score = self.player_b_score + new_player_b_pos

            new_state = GameState(new_turn_count,
                                  new_player_a_pos, new_player_a_score,
                                  new_player_b_pos, new_player_b_score)
            result[new_state] = occurrence_count
        return result

    def is_win_for_a(self):
        return self.player_a_score >= 21

    def is_win_for_b(self):
        return self.player_b_score >= 21

    def __repr__(self):
        return str(self)

    def __str__(self):
        if self.is_win_for_a():
            status = "A wins"
        elif self.is_win_for_b():
            status = "B wins"
        else:
            status = "%s to play" % ('A' if self.turn_count % 2 == 0 else 'B')

        return "%s: A at %d with %d, B at %d with %d, %s" % (
            self.turn_count,
            self.player_a_pos, self.player_a_score,
            self.player_b_pos, self.player_b_score,
            status)

    def __hash__(self):
        return hash(str(self))

    def __eq__(self, other):
        return str(self) == str(other)


def update_states(states):
    new_states = dict()
    for state, count in states.items():
        for new_state, new_count in state.next_states().items():
            new_states[new_state] = new_states.get(new_state, 0) + (count * new_count)
    return new_states


def play_nondetermininistic_to_win(player_a_start_pos, player_b_start_pos):
    initial_states = {GameState(0, player_a_start_pos, 0, player_b_start_pos, 0): 1}
    current_states = initial_states
    a_win_count = 0
    b_win_count = 0

    while len(current_states) > 0:
        new_states = update_states(current_states)
        to_remove = []

        for state, count in new_states.items():
            if state.is_win_for_a():
                a_win_count += count
                to_remove.append(state)
            elif state.is_win_for_b():
                b_win_count += count
                to_remove.append(state)

        for state in to_remove:
            del new_states[state]

        current_states = new_states

    return a_win_count, b_win_count


def test_play_nondetermininistic_to_win():
    assert play_nondetermininistic_to_win(4, 8) == (444356092776315, 341960390180808)
    print()
    print(max(play_nondetermininistic_to_win(7, 3)))