2018: d24: ex1: add solution

2018/d24/ex1/ex1.py

@@ -0,0 +1,188 @@
+#!/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()