aoc/year2021/
day18.rs

1//! # Snailfish
2//!
3//! The key observation is that snailfish numbers represent
4//! [binary trees](https://en.wikipedia.org/wiki/Binary_tree).
5//!
6//! For example the first four sample numbers on the problem description look like the following
7//! in binary tree form:
8//!
9//! ```text
10//! [1,2]    [[1,2],3]    [9,[8,7]]    [[1,9],[8,5]]
11//!   ■          ■            ■              ■
12//!  / \        / \          / \           /   \
13//! 1   2      ■   3        9   ■         ■     ■
14//!           / \              / \       / \   / \
15//!          1   2            8   7     1   9 8   5
16//! ```
17//!
18//! The addition rules have an important consequence. Exploding removes two leaf nodes at depth 5
19//! and moves them to neighbouring nodes. Since exploding repeatedly happens before splitting until
20//! there are no more values at depth 5 this means that the tree will never exceed a depth of 5.
21//!
22//! Each level of a tree can contain up to 2ⁿ nodes, so the maximum size of a snailfish tree is
23//! 1 + 2 + 4 + 8 + 16 + 32 = 2⁶-1 = 63 nodes.
24//!
25//! This means that we can store each snailfish number as an implicit data structure in a fixed size
26//! array. This is faster, smaller and more convenient than using a traditional struct with pointers.
27//! The root node is stored at index 0. For a node at index `i` its left child is at index
28//! `2i + 1`, right child at index `2i + 2` and parent at index `i / 2`. As leaf nodes are
29//! always greater than or equal to zero, `-1` is used as a special sentinel value for non-leaf nodes.
30use crate::util::parse::*;
31use crate::util::thread::*;
32use std::sync::atomic::{AtomicI32, Ordering};
33
34type Snailfish = [i32; 63];
35
36/// The indices for [in-order traversal](https://en.wikipedia.org/wiki/Tree_traversal) of the first
37/// 4 levels of the implicit binary tree stored in an array.
38const IN_ORDER: [usize; 30] = [
39    1, 3, 7, 15, 16, 8, 17, 18, 4, 9, 19, 20, 10, 21, 22, 2, 5, 11, 23, 24, 12, 25, 26, 6, 13, 27,
40    28, 14, 29, 30,
41];
42
43/// Parse a snailfish number into an implicit binary tree stored in an array.
44///
45/// Since no number will greater than 9 initially we can consider each character individually.
46/// `[` means moves down a level to parse children, `,` means move from left to right node,
47/// `]` means move up a level to return to parent and a digit from 0-9 creates a leaf node
48/// with that value.
49pub fn parse(input: &str) -> Vec<Snailfish> {
50    fn helper(bytes: &[u8]) -> Snailfish {
51        let mut tree = [-1; 63];
52        let mut i = 0;
53
54        for &b in bytes {
55            match b {
56                b'[' => i = 2 * i + 1,
57                b',' => i += 1,
58                b']' => i = (i - 1) / 2,
59                b => tree[i] = b.to_decimal() as i32,
60            }
61        }
62
63        tree
64    }
65    input.lines().map(|line| helper(line.as_bytes())).collect()
66}
67
68/// Add all snailfish numbers, reducing to a single magnitude.
69pub fn part1(input: &[Snailfish]) -> i32 {
70    let mut sum = input.iter().copied().reduce(|acc, n| add(&acc, &n)).unwrap();
71    magnitude(&mut sum)
72}
73
74/// Find the largest magnitude of any two snailfish numbers, remembering that snailfish addition
75/// is *not* commutative.
76pub fn part2(input: &[Snailfish]) -> i32 {
77    let mut pairs = Vec::new();
78
79    for (i, a) in input.iter().enumerate() {
80        for (j, b) in input.iter().enumerate() {
81            if i != j {
82                pairs.push((a, b));
83            }
84        }
85    }
86
87    // Use as many cores as possible to parallelize the calculation.
88    let shared = AtomicI32::new(0);
89    spawn_parallel_iterator(&pairs, |iter| worker(&shared, iter));
90    shared.into_inner()
91}
92
93/// Pair addition is independent so we can parallelize across multiple threads.
94fn worker(shared: &AtomicI32, iter: ParIter<'_, (&Snailfish, &Snailfish)>) {
95    let mut partial = 0;
96
97    for (a, b) in iter {
98        partial = partial.max(magnitude(&mut add(a, b)));
99    }
100
101    shared.fetch_max(partial, Ordering::Relaxed);
102}
103
104/// Add two snailfish numbers.
105///
106/// The initial step creates a new root node then makes the numbers the left and right children
107/// of this new root node, by copying the respective ranges of the implicit trees.
108///
109/// We can optimize the rules a little. This initial combination is the only time that more than one
110/// pair will be 4 levels deep simultaneously, so we can sweep from left to right on all possible
111/// leaf nodes in one pass.
112fn add(left: &Snailfish, right: &Snailfish) -> Snailfish {
113    let mut tree = [-1; 63];
114
115    tree[3..5].copy_from_slice(&left[1..3]);
116    tree[7..11].copy_from_slice(&left[3..7]);
117    tree[15..23].copy_from_slice(&left[7..15]);
118    tree[31..47].copy_from_slice(&left[15..31]);
119
120    tree[5..7].copy_from_slice(&right[1..3]);
121    tree[11..15].copy_from_slice(&right[3..7]);
122    tree[23..31].copy_from_slice(&right[7..15]);
123    tree[47..63].copy_from_slice(&right[15..31]);
124
125    for pair in (31..63).step_by(2) {
126        if tree[pair] >= 0 {
127            explode(&mut tree, pair);
128        }
129    }
130
131    while split(&mut tree) {}
132    tree
133}
134
135/// Explode a specific pair identified by an index.
136///
137/// Storing the tree as an implicit structure has a nice benefit that finding the next left or right
138/// node is straightforward. We first move to the next left or right leaf node by adding or
139/// subtracting one to the index. If this node is empty then we move to the parent node until we
140/// find a leaf node.
141///
142/// The leaf node at index 31 has no possible nodes to the left and similarly the leaf node at
143/// index 62 has no possible nodes to the right.
144fn explode(tree: &mut Snailfish, pair: usize) {
145    if pair > 31 {
146        let mut i = pair - 1;
147        loop {
148            if tree[i] >= 0 {
149                tree[i] += tree[pair];
150                break;
151            }
152            i = (i - 1) / 2;
153        }
154    }
155
156    if pair < 61 {
157        let mut i = pair + 2;
158        loop {
159            if tree[i] >= 0 {
160                tree[i] += tree[pair + 1];
161                break;
162            }
163            i = (i - 1) / 2;
164        }
165    }
166
167    tree[pair] = -1;
168    tree[pair + 1] = -1;
169    tree[(pair - 1) / 2] = 0;
170}
171
172/// Split a node into two child nodes.
173///
174/// Search the tree in an *in-order* traversal, splitting the first node over `10` found (if any).
175/// We can optimize the rules by immediately exploding if this results in a node 4 levels deep,
176/// as we know that the prior optimzation in the [`add`] function means that this is the only
177/// explosion possible.
178fn split(tree: &mut Snailfish) -> bool {
179    for &i in &IN_ORDER {
180        if tree[i] >= 10 {
181            tree[2 * i + 1] = tree[i] / 2;
182            tree[2 * i + 2] = (tree[i] + 1) / 2;
183            tree[i] = -1;
184
185            if i >= 15 {
186                explode(tree, 2 * i + 1);
187            }
188            return true;
189        }
190    }
191    false
192}
193
194/// Calculate the magnitude of a snailfish number in place without using recursion.
195///
196/// This operation is destructive but much faster than using a recursive approach and acceptable
197/// as we no longer need the original snailfish number afterwards.
198fn magnitude(tree: &mut Snailfish) -> i32 {
199    for i in (0..31).rev() {
200        if tree[i] == -1 {
201            tree[i] = 3 * tree[2 * i + 1] + 2 * tree[2 * i + 2];
202        }
203    }
204    tree[0]
205}