Pre-compute initial explode actions

src/year2021/day18.rs

@@ -24,9 +24,25 @@
 //!
 //! This means that we can store each snailfish number as an implicit data structure in a fixed-size
 //! array. This is faster, smaller and more convenient than using a traditional struct with pointers.
-//! The root node is stored at index 1 (index 0 is unused). For a node at index `i` its left child
-//! is at index `2i`, right child at index `2i + 1` and parent at index `i / 2`. As leaf nodes are
-//! always greater than or equal to zero, `-1` is used as a special sentinel value for non-leaf nodes.
+//! The root node is stored at index 1 (index 0 is unused by the tree, but see below). For a
+//! node at index `i` its left child is at index `2i`, right child at index `2i + 1` and parent
+//! at index `i / 2`. As leaf nodes are always greater than or equal to zero, `-1` is used as a
+//! special sentinel value for non-leaf nodes.
+//!
+//! Another optimization is realizing that all of the explode actions before the first split can
+//! be pre-computed. Instead of parsing a line `[1,[2,[3,[4,5]]]]` as written, we instead parse it
+//! it as if it had been `[[1,[2,[3,[4,5]]]],0]`, then perform the initial explode actions for that
+//! line up front, such that node 3 contains the value that must be added to the first leaf of a
+//! right-hand value in summation. Node 0 is not used by the implicit tree structure, so we instead
+//! use it as a tri-state value:
+//!
+//! - -2 means this Snailfish number is the result of a sum, for the left side during part one.
+//!   When passed to `add()`, we must shuffle contents one level lower.
+//! - -1 means this Snailfish number was just parsed, but did not explode left. When used on the
+//!   right side of add, any spillover from the left is added to the first leaf on the right.
+//! - Non-negative means this Snailfish number was just parsed, and had an explode that spilled
+//!   left. When used as the right side of an add, any spill from the left is combined with
+//!   this value, then added to the last leaf on the left.
 use crate::util::parse::*;
 use crate::util::thread::*;
 
@@ -49,8 +65,12 @@ pub fn parse(input: &str) -> Vec<Snailfish> {
     input
         .lines()
         .map(|line: &str| {
+            // Treat the line as if it had been `[line,0]`, then perform explode until it is back
+            // at depth 4. This allows later add() operations to do less work. Index 0 and 3
+            // then track the amount spilled left or right from those explodes.
             let mut tree = [-1; 64];
-            let mut i = 1;
+            tree[3] = 0;
+            let mut i = 2;
 
             for b in line.bytes() {
                 match b {
@@ -60,6 +80,11 @@ pub fn parse(input: &str) -> Vec<Snailfish> {
                     b => tree[i] = b.to_decimal() as i32,
                 }
             }
+            for pair in (32..48).step_by(2) {
+                if tree[pair] >= 0 {
+                    explode(&mut tree, pair);
+                }
+            }
 
             tree
         })
@@ -100,29 +125,56 @@ fn worker(iter: ParIter<'_, (&Snailfish, &Snailfish)>) -> Option<i32> {
 /// The initial step creates a new root node then makes the numbers the left and right children
 /// of this new root node, by copying the respective ranges of the implicit trees.
 ///
-/// We can optimize the rules a little. This initial combination is the only time that more than one
-/// pair will be 4 levels deep simultaneously, so we can sweep from left to right on all possible
-/// leaf nodes in one pass.
+/// We can optimize the rules a little. The parse step already ensured that there are no pairs
+/// deeper than 4 levels, and precomputed any explode values to spill between the two halves
+/// of the joined value. All that remains is checking for splits, where each split also takes
+/// care of any additional explodes needed.
 fn add(left: &Snailfish, right: &Snailfish) -> Snailfish {
     let mut tree = [-1; 64];
 
-    tree[4..6].copy_from_slice(&left[2..4]);
-    tree[8..12].copy_from_slice(&left[4..8]);
-    tree[16..24].copy_from_slice(&left[8..16]);
-    tree[32..48].copy_from_slice(&left[16..32]);
-
-    tree[6..8].copy_from_slice(&right[2..4]);
-    tree[12..16].copy_from_slice(&right[4..8]);
-    tree[24..32].copy_from_slice(&right[8..16]);
-    tree[48..64].copy_from_slice(&right[16..32]);
+    if left[0] == -2 {
+        // Left comes from a running sum during part one. We need to increase the depth, which
+        // in turn might cause some depth 5 leaves that need explode.
+        tree[3] = 0;
+        tree[4..6].copy_from_slice(&left[2..4]);
+        tree[8..12].copy_from_slice(&left[4..8]);
+        tree[16..24].copy_from_slice(&left[8..16]);
+        tree[32..48].copy_from_slice(&left[16..32]);
+
+        for pair in (32..48).step_by(2) {
+            if tree[pair] >= 0 {
+                explode(&mut tree, pair);
+            }
+        }
+    } else {
+        // We are adding two just-parsed numbers; the left is already rooted at 2 and has no depth 5
+        // leaves, making it ready to copy into place.
+        tree[3..24].copy_from_slice(&left[3..24]);
+    }
 
-    for pair in (32..64).step_by(2) {
-        if tree[pair] >= 0 {
-            explode(&mut tree, pair);
+    // Copy the right into place. This value is always just-parsed, with no depth 5 leaves.
+    tree[6..8].copy_from_slice(&right[4..6]);
+    tree[12..16].copy_from_slice(&right[8..12]);
+    tree[24..32].copy_from_slice(&right[16..24]);
+
+    // Adjust by the explode spillover between sides. We ensured that tree[3] contains any
+    // value to spill right, but must check right[0] to see if that sum then spills back left.
+    let (mut i, spill) = if right[0] == -1 { (24, tree[3]) } else { (23, tree[3] + right[0]) };
+    loop {
+        if tree[i] >= 0 {
+            tree[i] += spill;
+            break;
         }
+        i /= 2;
     }
+    tree[3] = -1;
 
+    // Now we process all split operations; any further explode actions are done during any split
+    // that creates a temporary depth 5.
     while split(&mut tree) {}
+
+    // Mark this tree as a sum before returning it.
+    tree[0] = -2;
     tree
 }
 
@@ -145,6 +197,9 @@ fn explode(tree: &mut Snailfish, pair: usize) {
             }
             i /= 2;
         }
+    } else {
+        // Store the left spill-out for later use by add().
+        tree[0] = tree[pair];
     }
 
     if pair < 62 {