aoc2015-22.py

# Day 22: Wizard Simulator 20XX
class GameState(NamedTuple):
    player: int
    mana: int
    boss: int
    damage: int
    shield: int = 0
    poison: int = 0
    recharge: int = 0


PRICES = {0: 53, 1: 73, 2: 113, 3: 173, 4: 229}


def player_round(state: GameState, spell: int, hard: bool) -> GameState:
    """
    0: Magic Missile costs 53 mana. It instantly does 4 damage.
    1: Drain costs 73 mana. It instantly does 2 damage and heals you for 2 hit points.
    2: Shield costs 113 mana. It starts an effect that lasts for 6 turns. While it is active, your armor is increased by 7.
    3: Poison costs 173 mana. It starts an effect that lasts for 6 turns. At the start of each turn while it is active, it deals the boss 3 damage.
    4: Recharge costs 229 mana. It starts an effect that lasts for 5 turns. At the start of each turn while it is active, it gives you 101 new mana.
    """
    player, mana, boss, damage, shield, poison, recharge = state
    if hard:
        player -= 1
        if player <= 0:
            # Player dies
            return GameState(player, mana, boss, damage, shield, poison, recharge)

    # Active effects
    if shield:
        # Shield only has effect during boss round, but ticks down here
        shield -= 1
    if poison:
        boss -= 3
        poison -= 1
    if recharge:
        mana += 101
        recharge -= 1

    # Spells
    # fmt: off
    match spell:
        case 0: boss -= 4
        case 1: player += 2; boss -= 2
        case 2: shield = 6
        case 3: poison = 6
        case 4: recharge = 5
    # fmt: on
    mana -= PRICES[spell]
    return GameState(player, mana, boss, damage, shield, poison, recharge)


def boss_round(state: GameState) -> GameState:
    "Return False if boss dies, otherwise new game state"
    player, mana, boss, damage, shield, poison, recharge = state
    # Active effects
    if shield:
        player += 7
        shield -= 1
    if poison:
        boss -= 3
        poison -= 1
        if boss <= 0:
            # Boss is dead, can't inflict any damage to player
            return GameState(player, mana, boss, damage, shield, poison, recharge)
    if recharge:
        mana += 101
        recharge -= 1
    player -= damage
    return GameState(player, mana, boss, damage, shield, poison, recharge)


def possible_spells(game: GameState) -> Set[GameState]:
    can_afford = set(spell for spell in PRICES if game.mana >= PRICES[spell])
    active_effects = set()
    if game.shield > 1:
        active_effects.add(2)
    if game.poison > 1:
        active_effects.add(3)
    if game.recharge > 1:
        active_effects.add(4)
    return can_afford - active_effects


@lru_cache(maxsize=None)
def play_rounds(game: GameState, hard: bool) -> int:

    best = inf

    for spell in possible_spells(game):
        state = player_round(game, spell, hard)
        if state.player > 0 and state.boss <= 0:
            return PRICES[spell]

        state = boss_round(state)
        if state.player > 0 and state.boss <= 0:
            return PRICES[spell]
        elif state.player <= 0:
            return inf

        new_cost = PRICES[spell] + play_rounds(state, hard)
        if new_cost < best:
            best = new_cost

    return best


state = GameState(player=50, mana=500, boss=55, damage=8)
print(f"Part 1: {play_rounds(state, hard=False)}")  # 953
print(f"Part 2: {play_rounds(state, hard=True)}")  # 1289