advent-of-code/2018/d24/ex1/ex1.py

#!/usr/bin/env python

import dataclasses
import enum
import sys


@dataclasses.dataclass
class Group:
    units: int
    hp: int
    weaknesses: set[str]
    immunities: set[str]
    attack: int
    attack_type: str
    initiative: int

    @classmethod
    def from_raw(cls, input: str) -> "Group":
        def split_sections(input: str) -> tuple[str, str, str]:
            points_idx = input.index("hit points ")
            with_idx = input.index(" with an attack")
            return (
                input[:points_idx].strip(),
                input[points_idx:with_idx].removeprefix("(").removesuffix(")"),
                input[with_idx:].strip(),
            )

        def parse_weak_immune(weak_immune: str) -> tuple[set[str], set[str]]:
            weaknesses: set[str] = set()
            immunities: set[str] = set()
            for part in weak_immune.split("; "):
                for start, values in (
                    ("weak to ", weaknesses),
                    ("immune to ", immunities),
                ):
                    if not part.startswith(start):
                        continue
                    values.update(part.removeprefix(start).split(", "))
            return weaknesses, immunities

        group_str, weak_immune, attack_str = split_sections(input)
        group_list, attack_list = group_str.split(), attack_str.split()
        weaknesses, immunities = parse_weak_immune(
            weak_immune.removeprefix("hit points (")
        )
        return cls(
            units=int(group_list[0]),
            hp=int(group_list[4]),
            weaknesses=weaknesses,
            immunities=immunities,
            attack=int(attack_list[5]),
            attack_type=attack_list[6],
            initiative=int(attack_list[10]),
        )

    @property
    def alive(self) -> bool:
        return self.units > 0

    @property
    def effective_power(self) -> int:
        return self.units * self.attack

    def potential_attack(self, ennemy: "Group") -> int:
        multiplier = 1
        if self.attack_type in ennemy.weaknesses:
            multiplier = 2
        if self.attack_type in ennemy.immunities:
            multiplier = 0
        return self.effective_power * multiplier


class Army(enum.StrEnum):
    INFECTION = "INFECTION"
    IMMUNE = "IMMUNE"

    def ennemy(self) -> "Army":
        if self == Army.INFECTION:
            return Army.IMMUNE
        if self == Army.IMMUNE:
            return Army.INFECTION
        assert False  # Sanity check


@dataclasses.dataclass
class Armies:
    immune: list[Group]
    infection: list[Group]

    @classmethod
    def from_raw(cls, input: str) -> "Armies":
        immune, infection = map(str.splitlines, input.split("\n\n"))
        assert "Immune System:" == immune[0]  # Sanity check
        assert "Infection:" == infection[0]  # Sanity check
        return cls(
            list(map(Group.from_raw, immune[1:])),
            list(map(Group.from_raw, infection[1:])),
        )

    def army(self, army: Army) -> list[Group]:
        if army == Army.IMMUNE:
            return self.immune
        if army == Army.INFECTION:
            return self.infection
        assert False  # Sanity check

    def active_groups(self, army: Army) -> set[int]:
        return {i for i, group in enumerate(self.army(army)) if group.alive}

    def selection_phase(self) -> dict[tuple[Army, int], int]:
        # Armies are sorted by decreasing power, initiative
        def power_order(group: Group) -> tuple[int, int]:
            return group.effective_power, group.initiative

        # Targets are ordered in decreasing potential attack, power, initiative
        def target_order(group: Group, ennemy: Group) -> tuple[int, int, int]:
            return (
                group.potential_attack(ennemy),
                ennemy.effective_power,
                ennemy.initiative,
            )

        res: dict[tuple[Army, int], int] = {}
        for army in Army:
            army_indices = sorted(
                self.active_groups(army),
                key=lambda i: power_order(self.army(army)[i]),
                reverse=True,
            )
            ennemies = self.army(army.ennemy())
            indices = set(self.active_groups(army.ennemy()))
            for i in army_indices:
                group = self.army(army)[i]
                if not indices:
                    break
                target = max(indices, key=lambda j: target_order(group, ennemies[j]))
                # Skip target if we cannot deal damage to it
                if group.potential_attack(ennemies[target]) == 0:
                    continue
                res[(army, i)] = target
                # Targets must be different for each attack
                indices.remove(target)
        return res

    def attack_phase(self, targets: dict[tuple[Army, int], int]) -> None:
        # Armies take turn by initiative, regardless of type
        turn_order = sorted(
            ((army, i) for army in Army for i in self.active_groups(army)),
            key=lambda t: self.army(t[0])[t[1]].initiative,
            reverse=True,
        )
        for army, i in turn_order:
            # Empty armies do not fight
            if not self.army(army)[i].alive:
                continue
            # Army must have a target selected
            if (target := targets.get((army, i))) is None:
                continue
            attackers = self.army(army)[i]
            defender = self.army(army.ennemy())[target]
            damage = attackers.potential_attack(defender)
            defender.units -= min(damage // defender.hp, defender.units)

    def fight(self) -> None:
        while self.active_groups(Army.IMMUNE) and self.active_groups(Army.INFECTION):
            targets = self.selection_phase()
            self.attack_phase(targets)


def solve(input: str) -> int:
    def parse(input: str) -> Armies:
        return Armies.from_raw(input)

    armies = parse(input)
    armies.fight()
    return sum(
        group.units for army in (armies.immune, armies.infection) for group in army
    )


def main() -> None:
    input = sys.stdin.read()
    print(solve(input))


if __name__ == "__main__":
    main()