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