Day 18 WIP, soo close.

2021/src/day18.rs

@@ -1,19 +1,35 @@
 use advent::prelude::*;
-use aoc_runner_derive::{aoc, aoc_generator};
+use aoc_runner_derive::aoc;
-use std::io::Read;
+use std::{
+    io::{BufReader, Cursor, Read},
+    ops::Add,
+};
 
-#[derive(Debug, PartialEq)]
+#[derive(Copy, Clone, Debug, PartialEq)]
 enum ChildType {
     None,
     Value(usize),
     Subtree(Idx),
 }
 
-#[derive(Copy, Clone, Debug, Default, PartialEq)]
+#[derive(Copy, Clone, Default, PartialEq)]
 struct Idx(usize);
 
+impl Debug for Idx {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
+impl Display for Idx {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
 #[derive(Debug, PartialEq)]
 struct Node {
+    deleted: bool,
     idx: Idx,
     parent: Option<Idx>,
     left: ChildType,
@@ -24,23 +40,197 @@ struct Node {
 // exploded spill over.
 // Need to support remove and/or replace for explode.
 // Need to support insert and/or replace for split.
-#[derive(Debug, Default, PartialEq)]
+#[derive(Debug, Default)]
 struct Tree {
     root: Idx,
     nodes: Vec<Node>,
 }
 
+struct TreeIter {
+    it: std::vec::IntoIter<Idx>,
+}
+
+impl TreeIter {
+    fn new(indices: &[Idx]) -> TreeIter {
+        let indices = indices.to_vec();
+        TreeIter {
+            it: indices.into_iter(),
+        }
+    }
+}
+
+impl Iterator for TreeIter {
+    type Item = Idx;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.it.next()
+    }
+}
+
+impl PartialEq for Tree {
+    fn eq(&self, other: &Self) -> bool {
+        // Lazy but should work.
+        self.to_string() == other.to_string()
+    }
+}
+
 fn read_byte<R: Read>(reader: &mut R) -> std::io::Result<Option<u8>> {
     reader.bytes().next().transpose()
 }
 
 impl Tree {
-    fn merge(left: &Tree, right: &Tree) -> Tree {
+    fn reduce(&mut self) {
-        // This is lazy but works for simple any obvious reasons (if FromStr and Display work
+        let mut changed = true;
-        // correctly).
+        while changed {
-        format!("[{},{}]", left, right)
+            changed = self.explode();
-            .parse()
+            println!("after explode {}", self);
-            .expect("failed to parse merge tree")
+            if changed {
+                continue;
+            }
+            changed = self.split();
+            println!("after split {}", self);
+            //println!("splice changed {}", changed);
+        }
+    }
+    fn magnitude(&self) -> usize {
+        fn inner(tree: &Tree, node: &Node) -> usize {
+            match (node.left, node.right) {
+                (ChildType::Value(l), ChildType::Value(r)) => 3 * l + 2 * r,
+                (ChildType::Subtree(idx), ChildType::Value(r)) => {
+                    3 * inner(&tree, &tree[idx]) + 2 * r
+                }
+                (ChildType::Value(l), ChildType::Subtree(idx)) => {
+                    3 * l + 2 * inner(&tree, &tree[idx])
+                }
+                (ChildType::Subtree(l_idx), ChildType::Subtree(r_idx)) => {
+                    3 * inner(&tree, &tree[l_idx]) + 2 * inner(&tree, &tree[r_idx])
+                }
+                _ => panic!("unhandled combo for magnitude"),
+            }
+        }
+        inner(self, &self[self.root])
+    }
+
+    fn split(&mut self) -> bool {
+        if let Some(split_idx) = self
+            .left_to_right()
+            .skip_while(|idx| {
+                let n = &self[*idx];
+                if let ChildType::Value(v) = n.left {
+                    if v > 9 {
+                        return false;
+                    }
+                }
+                if let ChildType::Value(v) = n.right {
+                    if v > 9 {
+                        return false;
+                    }
+                }
+                true
+            })
+            .next()
+        {
+            if let ChildType::Value(v) = self[split_idx].left {
+                if v > 9 {
+                    let l = v / 2;
+                    let r = if v % 2 == 1 { 1 + v / 2 } else { v / 2 };
+                    let mut new_idx = self.add_node(ChildType::Value(l), ChildType::Value(r));
+                    self[new_idx].parent = Some(split_idx);
+                    self[split_idx].left = ChildType::Subtree(new_idx);
+                }
+            }
+            if let ChildType::Value(v) = self[split_idx].right {
+                if v > 9 {
+                    let l = v / 2;
+                    let r = if v % 2 == 1 { 1 + v / 2 } else { v / 2 };
+                    let mut new_idx = self.add_node(ChildType::Value(l), ChildType::Value(r));
+                    self[new_idx].parent = Some(split_idx);
+                    self[split_idx].right = ChildType::Subtree(new_idx);
+                }
+            }
+            return true;
+        }
+        false
+    }
+    fn explode(&mut self) -> bool {
+        let mut changed = false;
+        if let Some(node) = self
+            .nodes
+            .iter()
+            .filter(|n| !n.deleted)
+            .find(|n| self.depth(n) >= 4)
+        {
+            changed = true;
+            let ex_idx = node.idx;
+            // Find spillover to the right
+            if let Some(spillover) = self
+                .left_to_right()
+                .skip_while(|idx| *idx != ex_idx)
+                .skip(1)
+                .find(|idx| {
+                    let n = &self[*idx];
+                    match (n.left, n.right) {
+                        (ChildType::Subtree(_), ChildType::Subtree(_)) => false,
+                        _ => true,
+                    }
+                })
+            {
+                let src = self[ex_idx].right;
+                let tgt = &mut self[spillover];
+                if let (ChildType::Value(l), ChildType::Value(r)) = (src, tgt.left) {
+                    tgt.left = ChildType::Value(l + r);
+                } else if let (ChildType::Value(l), ChildType::Value(r)) = (src, tgt.right) {
+                    tgt.right = ChildType::Value(l + r);
+                } else {
+                    unreachable!()
+                };
+            }
+            // Find spillover to the left
+            if let Some(spillover) = self
+                .right_to_left()
+                .skip_while(|idx| *idx != ex_idx)
+                .skip(1)
+                .find(|idx| {
+                    let n = &self[*idx];
+                    match (n.left, n.right) {
+                        (ChildType::Subtree(_), ChildType::Subtree(_)) => false,
+                        _ => true,
+                    }
+                })
+            {
+                let src = self[ex_idx].left;
+                let tgt = &mut self[spillover];
+                if let (ChildType::Value(l), ChildType::Value(r)) = (src, tgt.right) {
+                    tgt.right = ChildType::Value(l + r);
+                } else if let (ChildType::Value(l), ChildType::Value(r)) = (src, tgt.left) {
+                    tgt.left = ChildType::Value(l + r);
+                } else {
+                    unreachable!()
+                };
+            }
+            // Replace exploded node
+            self[ex_idx].deleted = true;
+            let p_idx = self[ex_idx].parent.expect("exploded root");
+            let p = &mut self[p_idx];
+            if let ChildType::Subtree(idx) = p.left {
+                if idx == ex_idx {
+                    p.left = ChildType::Value(0);
+                }
+            }
+            if let ChildType::Subtree(idx) = p.right {
+                if idx == ex_idx {
+                    p.right = ChildType::Value(0);
+                }
+            }
+        }
+        changed
+    }
+    fn depth(&self, node: &Node) -> usize {
+        if let Some(parent_idx) = node.parent {
+            1 + self.depth(&self[parent_idx])
+        } else {
+            0
+        }
     }
     fn find_root(&self, node: &Node) -> Idx {
         match node.parent {
@@ -51,6 +241,7 @@ impl Tree {
     fn add_node(&mut self, left: ChildType, right: ChildType) -> Idx {
         let idx = Idx(self.nodes.len());
         self.nodes.push(Node {
+            deleted: false,
             idx,
             parent: None,
             left,
@@ -59,7 +250,7 @@ impl Tree {
         idx
     }
 
-    fn from_str_node<R: Read>(&mut self, r: &mut R) -> ChildType {
+    fn from_str_node(&mut self, r: &mut BufReader<Cursor<&[u8]>>) -> ChildType {
         let mut parsing_left = true;
         // Can this be rewritten to eliminate the need for `None`?
         let mut left = ChildType::None;
@@ -94,7 +285,14 @@ impl Tree {
                 }
                 b',' => parsing_left = false,
                 b'0'..=b'9' => {
-                    let v = dbg!(b - b'0');
+                    let mut v = b - b'0';
+                    if let Ok(Some(peek)) = read_byte(r) {
+                        match peek {
+                            b'0'..=b'9' => v = (peek - b'0') + 10 * v,
+                            // Wasn't a number >9, push the byte back into the buffer.
+                            _ => r.seek_relative(-1).expect("failed to seek"),
+                        }
+                    }
                     if parsing_left {
                         left = ChildType::Value(v.into());
                         parsing_left = false;
@@ -113,18 +311,59 @@ impl Tree {
         write!(f, "[")?;
         match node.left {
             ChildType::None => panic!("left node was None"),
-            ChildType::Value(v) => write!(f, "{}", v),
+            ChildType::Value(v) => write!(f, "{}", v)?,
-            ChildType::Subtree(idx) => self.fmt_node(f, &self[idx]),
+            ChildType::Subtree(idx) => self.fmt_node(f, &self[idx])?,
         };
         write!(f, ",")?;
         match node.right {
             ChildType::None => panic!("right node was None"),
-            ChildType::Value(v) => write!(f, "{}", v),
+            ChildType::Value(v) => write!(f, "{}", v)?,
-            ChildType::Subtree(idx) => self.fmt_node(f, &self[idx]),
+            ChildType::Subtree(idx) => self.fmt_node(f, &self[idx])?,
         };
         write!(f, "]")?;
         Ok(())
     }
+
+    fn left_to_right(&mut self) -> TreeIter {
+        fn dfs(tree: &Tree, n: &Node, mut indices: &mut Vec<Idx>) {
+            if let ChildType::Subtree(idx) = n.left {
+                dfs(tree, &tree[idx], indices);
+            }
+            indices.push(n.idx);
+            if let ChildType::Subtree(idx) = n.right {
+                dfs(tree, &tree[idx], indices);
+            }
+        }
+        let mut indices = Vec::with_capacity(self.nodes.len());
+        dfs(self, &self[self.root], &mut indices);
+        TreeIter::new(&indices)
+    }
+    fn right_to_left(&mut self) -> TreeIter {
+        fn dfs(tree: &Tree, n: &Node, mut indices: &mut Vec<Idx>) {
+            if let ChildType::Subtree(idx) = n.right {
+                dfs(tree, &tree[idx], indices);
+            }
+            indices.push(n.idx);
+            if let ChildType::Subtree(idx) = n.left {
+                dfs(tree, &tree[idx], indices);
+            }
+        }
+        let mut indices = Vec::with_capacity(self.nodes.len());
+        dfs(self, &self[self.root], &mut indices);
+        TreeIter::new(&indices)
+    }
+}
+
+impl Add for Tree {
+    type Output = Tree;
+
+    fn add(self, other: Self) -> Self {
+        // This is lazy but works for simple any obvious reasons (if FromStr and Display work
+        // correctly).
+        format!("[{},{}]", self, other)
+            .parse()
+            .expect("failed to parse merge tree")
+    }
 }
 
 impl FromStr for Tree {
@@ -137,7 +376,8 @@ impl FromStr for Tree {
             read_byte(&mut bytes).expect("couldn't read first byte"),
             Some(b'[')
         );
-        tree.from_str_node(&mut bytes);
+        let mut b = BufReader::new(Cursor::new(bytes));
+        tree.from_str_node(&mut b);
         tree.root = tree.find_root(&tree[Idx(0)]);
         Ok(tree)
     }
@@ -168,13 +408,34 @@ impl Display for Tree {
     }
 }
 
+fn sum(input: &str) -> Tree {
+    input
+        .lines()
+        .map(|l| l.parse().expect("failed to parse"))
+        .reduce(|acc, t| acc + t)
+        .expect("failed to reduce")
+}
+
 #[aoc(day18, part1)]
 fn part1(input: &str) -> Result<usize> {
-    for l in input.lines() {
+    let nums: Vec<Tree> = input
-        let tree: Tree = l.parse()?;
+        .lines()
-        dbg!(&tree);
+        .map(|l| {
+            dbg!(l);
+            l.parse().expect("failed to parse")
+        })
+        .collect();
+    let mut it = nums.into_iter();
+    let mut last = it.next().unwrap();
+    while let Some(next) = it.next() {
+        println!("  {}", last);
+        println!("+ {}", next);
+        last = last + next;
+        println!("= {}", last);
+        last.reduce();
+        println!("= {}\n", last);
     }
-    Ok(0)
+    Ok(last.magnitude())
 }
 
 /*
@@ -203,21 +464,244 @@ mod tests {
     }
 
     #[test]
-    fn test_merge() -> Result<()> {
+    fn test_sum() -> Result<()> {
-        assert_eq!(
+        let l: Tree = "[1,2]".parse().unwrap();
-            Tree::merge(&"[1,2]".parse().unwrap(), &"[[3,4],5]".parse().unwrap()),
+        let r: Tree = "[[3,4],5]".parse().unwrap();
-            "[[1,2],[[3,4],5]]".parse().unwrap()
+        assert_eq!(l + r, "[[1,2],[[3,4],5]]".parse().unwrap());
-        );
+
+        let input = r#"
+[[[[4,3],4],4],[7,[[8,4],9]]]
+[1,1]
+    "#
+        .trim();
+        let mut s = sum(input);
+        s.reduce();
+        assert_eq!(s.to_string(), "[[[[0,7],4],[[7,8],[6,0]]],[8,1]]");
         Ok(())
     }
 
-    //#[test]
+    #[test]
+    fn test_reduce() -> Result<()> {
+        for (input, want) in [
+            ("[0,[0,[0,[0,[0,0]]]]]", "[0,[0,[0,[0,0]]]]"),
+            ("[[[[[[[[0,0],0],0],0],0],0],0],0]", "[[[[0,0],0],0],0]"),
+            ("[[[[[[[0,0],0],0],0],0],0],0]", "[[[[0,0],0],0],0]"),
+        ] {
+            println!("== test_reduce: {}", input);
+            let mut tree: Tree = input.parse()?;
+            tree.reduce();
+            let want = want.parse()?;
+            assert_eq!(tree, want, "\nInput {} Got {} Want {}", input, tree, want);
+        }
+        Ok(())
+    }
+    #[test]
+    fn test_explode() -> Result<()> {
+        for (input, want) in [
+            ("[[[[0,0],0],0],0]", "[[[[0,0],0],0],0]"),
+            ("[[[0,0],0],0]", "[[[0,0],0],0]"),
+            ("[[0,0],0]", "[[0,0],0]"),
+            ("[0,[0,[0,[0,0]]]]", "[0,[0,[0,[0,0]]]]"),
+            ("[0,[0,[0,0]]]", "[0,[0,[0,0]]]"),
+            ("[0,[0,0]]", "[0,[0,0]]"),
+            ("[[[[[9,8],1],2],3],4]", "[[[[0,9],2],3],4]"),
+            ("[7,[6,[5,[4,[3,2]]]]]", "[7,[6,[5,[7,0]]]]"),
+            ("[[6,[5,[4,[3,2]]]],1]", "[[6,[5,[7,0]]],3]"),
+            (
+                "[[3,[2,[1,[7,3]]]],[6,[5,[4,[3,2]]]]]",
+                "[[3,[2,[8,0]]],[9,[5,[4,[3,2]]]]]",
+            ),
+            (
+                "[[3,[2,[8,0]]],[9,[5,[4,[3,2]]]]]",
+                "[[3,[2,[8,0]]],[9,[5,[7,0]]]]",
+            ),
+        ] {
+            println!("== test_explode: {}", input);
+            let mut tree: Tree = input.parse()?;
+            tree.explode();
+            let want = want.parse()?;
+            assert_eq!(tree, want, "\nInput {} Got {} Want {}", input, tree, want);
+        }
+        Ok(())
+    }
+
+    #[test]
+    fn test_split() -> Result<()> {
+        for (input, want) in [
+            ("[10,0]", "[[5,5],0]"), //
+            ("[0,11]", "[0,[5,6]]"),
+            ("[[0,11],0]", "[[0,[5,6]],0]"),
+            ("[11,0]", "[[5,6],0]"),
+            ("[0,[11,0]]", "[0,[[5,6],0]]"),
+            ("[12,0]", "[[6,6],0]"),
+            ("[0,12]", "[0,[6,6]]"),
+        ] {
+            println!("== test_split: {}", input);
+            let mut tree: Tree = input.parse()?;
+            dbg!(&tree);
+            tree.split();
+            let want = want.parse()?;
+            assert_eq!(tree, want, "\nInput {} Got {} Want {}", input, tree, want);
+        }
+        Ok(())
+    }
+
+    #[test]
+    fn test_magnitude() -> Result<()> {
+        for (input, want) in [
+            ("[9,1]", 29),
+            ("[1,9]", 21),
+            ("[[9,1],[1,9]]", 129),
+            ("[[1,2],[[3,4],5]]", 143),
+            ("[[[[0,7],4],[[7,8],[6,0]]],[8,1]]", 1384),
+            ("[[[[1,1],[2,2]],[3,3]],[4,4]]", 445),
+            ("[[[[3,0],[5,3]],[4,4]],[5,5]]", 791),
+            ("[[[[5,0],[7,4]],[5,5]],[6,6]]", 1137),
+            (
+                "[[[[8,7],[7,7]],[[8,6],[7,7]]],[[[0,7],[6,6]],[8,7]]]",
+                3488,
+            ),
+        ] {
+            let tree: Tree = input.parse()?;
+            assert_eq!(tree.magnitude(), want);
+        }
+        Ok(())
+    }
+
+    #[test]
+    fn test_add_and_reduce() -> Result<()> {
+        for (input, want) in [
+            (
+                r#"
+[1,1]
+[2,2]
+[3,3]
+[4,4]
+"#,
+                "[[[[1,1],[2,2]],[3,3]],[4,4]]",
+            ),
+            (
+                r#"
+[1,1]
+[2,2]
+[3,3]
+[4,4]
+[5,5]
+    "#
+                .trim(),
+                "[[[[3,0],[5,3]],[4,4]],[5,5]]",
+            ),
+            (
+                r#"
+[1,1]
+[2,2]
+[3,3]
+[4,4]
+[5,5]
+[6,6]
+    "#
+                .trim(),
+                "[[[[5,0],[7,4]],[5,5]],[6,6]]",
+            ),
+            (
+                r#"
+[[[[4,3],4],4],[7,[[8,4],9]]]
+[1,1]
+"#
+                .trim(),
+                "[[[[0,7],4],[[7,8],[6,0]]],[8,1]]",
+            ),
+        ] {
+            println!("== 1. test_add_and_reduce: {}", input);
+            let mut num = sum(input.trim());
+            println!("before reduce: {}", num);
+            num.reduce();
+            println!("after reduce: {}", num);
+            assert_eq!(num.to_string(), want);
+        }
+
+        for (l, r, eq) in [
+            (
+                "[[[[4,3],4],4],[7,[[8,4],9]]]",
+                "[1,1]",
+                "[[[[0,7],4],[[7,8],[6,0]]],[8,1]]",
+            ),
+            (
+                "[[[0,[4,5]],[0,0]],[[[4,5],[2,6]],[9,5]]]",
+                "[7,[[[3,7],[4,3]],[[6,3],[8,8]]]]",
+                "[[[[4,0],[5,4]],[[7,7],[6,0]]],[[8,[7,7]],[[7,9],[5,0]]]]",
+            ),
+            (
+                "[[[[4,0],[5,4]],[[7,7],[6,0]]],[[8,[7,7]],[[7,9],[5,0]]]]",
+                "[[2,[[0,8],[3,4]]],[[[6,7],1],[7,[1,6]]]]",
+                "[[[[6,7],[6,7]],[[7,7],[0,7]]],[[[8,7],[7,7]],[[8,8],[8,0]]]]",
+            ),
+            (
+                "[[[[6,7],[6,7]],[[7,7],[0,7]]],[[[8,7],[7,7]],[[8,8],[8,0]]]]",
+                "[[[[2,4],7],[6,[0,5]]],[[[6,8],[2,8]],[[2,1],[4,5]]]]",
+                "[[[[7,0],[7,7]],[[7,7],[7,8]]],[[[7,7],[8,8]],[[7,7],[8,7]]]]",
+            ),
+            (
+                "[[[[7,0],[7,7]],[[7,7],[7,8]]],[[[7,7],[8,8]],[[7,7],[8,7]]]]",
+                "[7,[5,[[3,8],[1,4]]]]",
+                "[[[[7,7],[7,8]],[[9,5],[8,7]]],[[[6,8],[0,8]],[[9,9],[9,0]]]]",
+            ),
+            (
+                "[[[[7,7],[7,8]],[[9,5],[8,7]]],[[[6,8],[0,8]],[[9,9],[9,0]]]]",
+                "[[2,[2,2]],[8,[8,1]]]",
+                "[[[[6,6],[6,6]],[[6,0],[6,7]]],[[[7,7],[8,9]],[8,[8,1]]]]",
+            ),
+            (
+                "[[[[6,6],[6,6]],[[6,0],[6,7]]],[[[7,7],[8,9]],[8,[8,1]]]]",
+                "[2,9]",
+                "[[[[6,6],[7,7]],[[0,7],[7,7]]],[[[5,5],[5,6]],9]]",
+            ),
+            (
+                "[[[[6,6],[7,7]],[[0,7],[7,7]]],[[[5,5],[5,6]],9]]",
+                "[1,[[[9,3],9],[[9,0],[0,7]]]]",
+                "[[[[7,8],[6,7]],[[6,8],[0,8]]],[[[7,7],[5,0]],[[5,5],[5,6]]]]",
+            ),
+            (
+                "[[[[7,8],[6,7]],[[6,8],[0,8]]],[[[7,7],[5,0]],[[5,5],[5,6]]]]",
+                "[[[5,[7,4]],7],1]",
+                "[[[[7,7],[7,7]],[[8,7],[8,7]]],[[[7,0],[7,7]],9]]",
+            ),
+            (
+                "[[[[7,7],[7,7]],[[8,7],[8,7]]],[[[7,0],[7,7]],9]]",
+                "[[[[4,2],2],6],[8,7]]",
+                "[[[[8,7],[7,7]],[[8,6],[7,7]]],[[[0,7],[6,6]],[8,7]]]",
+            ),
+        ] {
+            let l: Tree = l.parse()?;
+            let r: Tree = r.parse()?;
+            let mut num = l + r;
+            println!("== 2. test_add_and_reduce: {}", num);
+            println!("before reduce: {}", num);
+            num.reduce();
+            println!("after reduce: {}", num);
+            assert_eq!(num.to_string(), eq);
+        }
+        Ok(())
+    }
+
+    #[test]
     fn test_part1() -> Result<()> {
-        let sum = [
+        let input = r#"
-            [[[6, 6], [7, 6]], [[7, 7], [7, 0]]],
+[[[0,[4,5]],[0,0]],[[[4,5],[2,6]],[9,5]]]
-            [[[7, 7], [7, 7]], [[7, 8], [9, 9]]],
+[7,[[[3,7],[4,3]],[[6,3],[8,8]]]]
-        ];
+[[2,[[0,8],[3,4]]],[[[6,7],1],[7,[1,6]]]]
-        dbg!(&sum);
+[[[[2,4],7],[6,[0,5]]],[[[6,8],[2,8]],[[2,1],[4,5]]]]
+[7,[5,[[3,8],[1,4]]]]
+[[2,[2,2]],[8,[8,1]]]
+[2,9]
+[1,[[[9,3],9],[[9,0],[0,7]]]]
+[[[5,[7,4]],7],1]
+[[[[4,2],2],6],[8,7]]
+    "#
+        .trim();
+
+        assert_eq!(part1(input)?, 3488);
+
         let input = r#"
 [[[0,[5,8]],[[1,7],[9,6]]],[[4,[1,2]],[[1,4],2]]]
 [[[5,[2,8]],4],[5,[[9,9],0]]]