From fc803d8a0f89f00f0a8e8b013c9928adc488159a Mon Sep 17 00:00:00 2001
From: Johannes <freelon@gmx.net>
Date: Tue, 14 Nov 2023 19:35:00 +0100
Subject: [PATCH] Day 24 part 2.

---
 src/main.rs        |  2 +-
 src/tasks/day24.rs | 76 +++++++++++++++++++++++++++++++++++++---------
 2 files changed, 62 insertions(+), 16 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index e0d66ee..f9c8efb 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,3 +1,3 @@
 fn main() {
-    aoc_2018::tasks::day24::task1();
+    aoc_2018::tasks::day24::task2();
 }
diff --git a/src/tasks/day24.rs b/src/tasks/day24.rs
index cd6faee..97ba547 100644
--- a/src/tasks/day24.rs
+++ b/src/tasks/day24.rs
@@ -1,17 +1,32 @@
 use std::cmp::Reverse;
 use std::collections::HashMap;
 
+use itertools::Itertools;
+
 use crate::utils;
 
 pub fn task1() {
     let input = utils::read_file("input/day24.txt");
 
-    let remaining_units = play(&input);
-    println!("Standing units after the game: {}", remaining_units);
+    let result = play(&input, 0);
+    println!("Standing units after the game: {}", result.unwrap().1);
 }
 
-fn play(input: &str) -> i64 {
-    let mut game = Game::create(&input);
+pub fn task2() {
+    let input = utils::read_file("input/day24.txt");
+
+    let immune_final_count = (0..)
+        .map(|boost| play(&input, boost))
+        .find(|result| result.map(|(team, _)| team) == Some(Team::ImmuneSystem))
+        .unwrap()
+        .unwrap()
+        .1;
+
+    println!("Immune systems final count: {immune_final_count}");
+}
+
+fn play(input: &str, boost: i64) -> Option<(Team, i64)> {
+    let mut game = Game::create(&input, boost);
 
     println!(
         "Immune start units: {}",
@@ -30,18 +45,20 @@ fn play(input: &str) -> i64 {
             .sum::<i64>()
     );
 
-    let mut rounds_played = 0;
     while !game.is_over() {
         game.round();
-        rounds_played += 1;
+        if game.draw {
+            return None;
+        }
     }
 
-    game.remaining_units()
+    Some((game.winning_team(), game.remaining_units()))
 }
 
 #[derive(Debug)]
 struct Game {
     groups: HashMap<usize, Group>,
+    draw: bool,
 }
 
 impl Game {
@@ -49,7 +66,20 @@ impl Game {
         self.groups.values().map(|it| it.units).sum::<i64>()
     }
 
-    fn create(input: &str) -> Self {
+    fn winning_team(&self) -> Team {
+        let teams: Vec<Team> = self
+            .groups
+            .values()
+            .map(|group| group.team)
+            .unique()
+            .collect();
+        if teams.len() != 1 {
+            panic!("No winning team. Remaining teams: {teams:?}");
+        }
+        teams[0]
+    }
+
+    fn create(input: &str, boost: i64) -> Self {
         let mut groups = HashMap::new();
         let mut team = Team::ImmuneSystem;
         for (id, line) in input.lines().enumerate() {
@@ -62,7 +92,10 @@ impl Game {
                 }
                 "" => (),
                 group => {
-                    if let Some(group) = Group::from_str(group, team, id) {
+                    if let Some(mut group) = Group::from_str(group, team, id) {
+                        if group.team == Team::ImmuneSystem {
+                            group.attack_damage += boost;
+                        }
                         groups.insert(id, group);
                     } else {
                         panic!("bad group: {group}");
@@ -70,7 +103,10 @@ impl Game {
                 }
             }
         }
-        Game { groups }
+        Game {
+            groups,
+            draw: false,
+        }
     }
 
     fn round(&mut self) {
@@ -125,6 +161,7 @@ impl Game {
         let mut all_ids_by_initiative: Vec<usize> = self.groups.values().map(|it| it.id).collect();
         all_ids_by_initiative.sort_unstable_by_key(|id| Reverse(self.groups[id].initiative));
 
+        self.draw = true;
         for active_id in all_ids_by_initiative {
             if !self.groups[&active_id].alive() {
                 // was killed in this round
@@ -137,6 +174,9 @@ impl Game {
                 let dying_units = damage / enemy.hp_each;
                 self.groups.get_mut(enemy_id).unwrap().units -= dying_units;
 
+                if dying_units > 0 {
+                    self.draw = false;
+                }
                 // println!(
                 //     "{} dealt {} ({} units) damage to {}",
                 //     active_id, damage, dying_units, enemy_id
@@ -158,7 +198,7 @@ impl Game {
     }
 }
 
-#[derive(Debug, Copy, Clone, PartialEq)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
 enum Team {
     Infection,
     ImmuneSystem,
@@ -247,15 +287,21 @@ impl Group {
 mod test {
     use crate::tasks::day24::play;
 
-    #[test]
-    fn example1() {
-        let input = "Immune System:
+    const INPUT: &str = "Immune System:
     17 units each with 5390 hit points (weak to radiation, bludgeoning) with an attack that does 4507 fire damage at initiative 2
     989 units each with 1274 hit points (immune to fire; weak to bludgeoning, slashing) with an attack that does 25 slashing damage at initiative 3
 
 Infection:
     801 units each with 4706 hit points (weak to radiation) with an attack that does 116 bludgeoning damage at initiative 1
     4485 units each with 2961 hit points (immune to radiation; weak to fire, cold) with an attack that does 12 slashing damage at initiative 4";
-        assert_eq!(5216, play(input));
+
+    #[test]
+    fn example1() {
+        assert_eq!(5216, play(INPUT, 0).unwrap().1);
+    }
+
+    #[test]
+    fn example2() {
+        assert_eq!(51, play(INPUT, 1570).unwrap().1);
     }
 }