From 47cb2fefa5a44424f8f70f3900972fb43938957b Mon Sep 17 00:00:00 2001
From: jazzpi <jasper@mezzo.de>
Date: Wed, 21 Dec 2022 18:45:28 +0100
Subject: [PATCH] Day 21, part 2

---
 src/bin/d21p1.rs |  2 +-
 src/bin/d21p2.rs |  9 +++++
 src/day21.rs     | 98 ++++++++++++++++++++++++++++++++++++++++--------
 3 files changed, 92 insertions(+), 17 deletions(-)
 create mode 100644 src/bin/d21p2.rs

diff --git a/src/bin/d21p1.rs b/src/bin/d21p1.rs
index bd9ff0f..ec9d46e 100644
--- a/src/bin/d21p1.rs
+++ b/src/bin/d21p1.rs
@@ -3,5 +3,5 @@ use aoc22::{day21, util};
 pub fn main() {
     let monkeys = day21::parse_monkeys(&util::parse_input());
 
-    println!("root yells: {}", day21::yell(monkeys.root()));
+    println!("root yells: {}", day21::yell(monkeys.root(), false));
 }
diff --git a/src/bin/d21p2.rs b/src/bin/d21p2.rs
new file mode 100644
index 0000000..eb19999
--- /dev/null
+++ b/src/bin/d21p2.rs
@@ -0,0 +1,9 @@
+use aoc22::{day21, util};
+
+pub fn main() {
+    let monkeys = day21::parse_monkeys(&util::parse_input());
+
+    let human = day21::find(&monkeys, "humn").unwrap();
+
+    println!("Human should yell: {}", day21::find_solution(human));
+}
diff --git a/src/day21.rs b/src/day21.rs
index e2a4004..2cf2a0b 100644
--- a/src/day21.rs
+++ b/src/day21.rs
@@ -3,7 +3,7 @@ use std::collections::HashMap;
 use regex::Regex;
 use trees::{tr, Node, Tree};
 
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy, PartialEq)]
 pub enum Operation {
     Add,
     Sub,
@@ -23,12 +23,18 @@ impl Operation {
     }
 }
 
-#[derive(Debug, Clone)]
-pub enum Monkey {
+#[derive(Debug, Clone, PartialEq)]
+pub enum MonkeyKind {
     Number(isize),
     Calculate(String, String, Operation),
 }
 
+#[derive(Debug, Clone, PartialEq)]
+pub struct Monkey {
+    pub name: String,
+    pub kind: MonkeyKind,
+}
+
 pub fn parse_monkeys(input: &String) -> Tree<Monkey> {
     let mut monkeys = HashMap::new();
 
@@ -36,16 +42,22 @@ pub fn parse_monkeys(input: &String) -> Tree<Monkey> {
     for line in input.lines() {
         let captures = re.captures(line).unwrap();
         let name = captures.get(1).unwrap().as_str();
-        let monkey = if let Some(num) = captures.get(2) {
-            Monkey::Number(num.as_str().parse().unwrap())
+        let kind = if let Some(num) = captures.get(2) {
+            MonkeyKind::Number(num.as_str().parse().unwrap())
         } else {
-            Monkey::Calculate(
+            MonkeyKind::Calculate(
                 captures.get(3).unwrap().as_str().to_owned(),
                 captures.get(5).unwrap().as_str().to_owned(),
                 Operation::from(captures.get(4).unwrap().as_str()),
             )
         };
-        monkeys.insert(name.to_owned(), monkey);
+        monkeys.insert(
+            name.to_owned(),
+            Monkey {
+                name: name.to_owned(),
+                kind,
+            },
+        );
     }
 
     let tree = to_tree(&monkeys, "root");
@@ -57,7 +69,7 @@ fn to_tree(monkeys: &HashMap<String, Monkey>, monkey: &str) -> Tree<Monkey> {
     let monkey = &monkeys[monkey];
     let mut tree = tr(monkey.clone());
 
-    if let Monkey::Calculate(a, b, _) = monkey {
+    if let MonkeyKind::Calculate(a, b, _) = &monkey.kind {
         let forest = -to_tree(monkeys, a) - to_tree(monkeys, b);
         tree.append(forest);
     }
@@ -65,19 +77,73 @@ fn to_tree(monkeys: &HashMap<String, Monkey>, monkey: &str) -> Tree<Monkey> {
     tree
 }
 
-pub fn yell(monkey: &Node<Monkey>) -> isize {
-    match monkey.data() {
-        Monkey::Number(n) => *n,
-        Monkey::Calculate(_, _, op) => {
+pub fn yell(monkey: &Node<Monkey>, no_human: bool) -> isize {
+    if no_human {
+        assert!(monkey.data().name != "humn");
+    }
+
+    match &monkey.data().kind {
+        MonkeyKind::Number(n) => *n,
+        MonkeyKind::Calculate(_, _, op) => {
             let mut children = monkey.iter();
             let a = children.next().unwrap();
             let b = children.next().unwrap();
             match op {
-                Operation::Add => yell(a) + yell(b),
-                Operation::Sub => yell(a) - yell(b),
-                Operation::Mul => yell(a) * yell(b),
-                Operation::Div => yell(a) / yell(b),
+                Operation::Add => yell(a, no_human) + yell(b, no_human),
+                Operation::Sub => yell(a, no_human) - yell(b, no_human),
+                Operation::Mul => yell(a, no_human) * yell(b, no_human),
+                Operation::Div => yell(a, no_human) / yell(b, no_human),
             }
         }
     }
 }
+
+pub fn find_solution(target: &Node<Monkey>) -> isize {
+    let parent = target.parent().unwrap();
+    let other_pos = parent.iter().position(|n| n != target).unwrap();
+    assert!(other_pos < 2);
+    let sibling = parent.iter().nth(other_pos).unwrap();
+    let other = yell(sibling, true);
+
+    if parent.parent().is_none() {
+        other
+    } else {
+        let parent_target = find_solution(parent);
+        if let MonkeyKind::Calculate(_, _, op) = parent.data().kind {
+            match op {
+                Operation::Add => parent_target - other,
+                Operation::Sub => {
+                    if other_pos == 0 {
+                        other - parent_target
+                    } else {
+                        other + parent_target
+                    }
+                }
+                Operation::Mul => parent_target / other,
+                Operation::Div => {
+                    if other_pos == 0 {
+                        other / parent_target
+                    } else {
+                        parent_target * other
+                    }
+                }
+            }
+        } else {
+            panic!("Parent does not calculate!?");
+        }
+    }
+}
+
+pub fn find<'a>(monkeys: &'a Tree<Monkey>, name: &str) -> Option<&'a Node<Monkey>> {
+    let mut next = vec![monkeys.root()];
+
+    while let Some(n) = next.pop() {
+        let data = n.data();
+        if data.name == name {
+            return Some(n);
+        }
+        next.extend(n.iter());
+    }
+
+    None
+}