Skip to main content

aoc/year2021/
day03.rs

1//! # Binary Diagnostic
2//!
3//! Part one collects 12 frequency counters, one for each digit lane. Each line of input increments
4//! the corresponding counters for each `1`. The number of `0` in the lane can be recovered by
5//! subtraction, making it easy to construct both the gamma and epsilon rate.
6//!
7//! Part two constructs a binary tree, stored as an array. Using an array of 2*4096 bins is enough
8//! to treat each of the 1000 unique inputs as a leaf in the second half of the array, and then
9//! each node in the first half is computed as the sum of its two children from the right half. The
10//! array can be built in linear time, at which point finding the oxygen generator rating and
11//! CO2 scrubber rating are each a binary search in the tree.
12pub struct Input {
13    size: usize,       // Total number of binary digits in each number.
14    entries: usize,    // Number of entries in the file.
15    lanes: Vec<usize>, // Contains size counters, one for each digit lane.
16    tree: Vec<usize>,  // Contains a balanced binary tree with 2**(width+1) nodes.
17}
18
19pub fn parse(input: &str) -> Input {
20    let size = input.lines().next().unwrap().len();
21    let offset = 1 << size;
22    let entries = input.len() / (size + 1);
23
24    // Allocate frequency counters based on size of each input.
25    let mut lanes = vec![0; size];
26    let mut tree = vec![0; 1 << (size + 1)];
27
28    // Update frequency counters and leaf nodes for each number in the input.
29    for chunk in input.as_bytes().chunks(size + 1) {
30        let mut binary = 0;
31
32        for (lane, b) in lanes.iter_mut().zip(&chunk[..size]) {
33            let bit = usize::from(b & 1);
34            *lane += bit;
35            binary = (binary << 1) | bit;
36        }
37
38        tree[offset + binary] = 1;
39    }
40
41    // Finish populating children counts in the tree.
42    for i in (1..offset).rev() {
43        tree[i] = tree[2 * i] + tree[2 * i + 1];
44    }
45
46    Input { size, entries, lanes, tree }
47}
48
49pub fn part1(input: &Input) -> u32 {
50    let (gamma, epsilon) = input.lanes.iter().fold((0, 0), |(gamma, epsilon), &ones| {
51        let zeros = input.entries - ones;
52        ((gamma << 1) | u32::from(ones > zeros), (epsilon << 1) | u32::from(zeros > ones))
53    });
54
55    gamma * epsilon
56}
57
58pub fn part2(input: &Input) -> usize {
59    let Input { size, tree, .. } = input;
60    let offset = 1 << size;
61
62    let mut ogr = 2; // Oxygen generator rating.
63    let mut csr = 2; // CO2 scrubber rating.
64
65    // Perform a binary search over the tree.
66    for _ in 0..*size {
67        ogr = 2 * (ogr + usize::from(tree[ogr + 1] >= tree[ogr]));
68        csr = 2 * (csr + usize::from(tree[csr + 1].wrapping_sub(1) < tree[csr].wrapping_sub(1)));
69    }
70
71    (ogr / 2 - offset) * (csr / 2 - offset)
72}