Advent of Code day 18 in decent time complexity

main.rs

use itertools::Itertools;
use std::collections::HashSet;

const N6:  [(i64, i64, i64); 6] = [
    ( 1, 0, 0), (0,  1, 0), (0, 0,  1),
    (-1, 0, 0), (0, -1, 0), (0, 0, -1)
];
const N14: [(i64, i64, i64); 18]= [
    ( 1,  0,  0), (0,  1,  0), ( 0,  0,  1),
    (-1,  0,  0), (0, -1,  0), ( 0,  0, -1),
    (-1,  1,  0), (0, -1,  1), ( 1,  0, -1),
    ( 1, -1,  0), (0,  1, -1), (-1,  0,  1),
    ( 1,  1,  0), (0,  1,  1), ( 1,  0,  1),
    (-1, -1,  0), (0, -1, -1), (-1,  0, -1)
];

fn displace(point: (i64, i64, i64), offsets: &[(i64, i64, i64)]) -> impl Iterator<Item=(i64, i64, i64)> + '_ {
    offsets.iter().map(move |off| (point.0 + off.0, point.1 + off.1, point.2 + off.2))
}

fn connected_component_of(points: &mut HashSet<(i64, i64, i64)>, start: (i64, i64, i64), neighbors: &[(i64, i64, i64)]) -> HashSet<(i64, i64, i64)> {
    let mut c = HashSet::new();
    let mut b = vec![start];
    points.remove(&start);

    while let Some(it) = b.pop() {
        assert!(!points.contains(&it));
        points.remove(&it);
        c.insert(it);
        for p in displace(it, neighbors) {
            if points.contains(&p) {
                points.remove(&p);
                b.push(p);
            }
        }
    }
    c
}

fn connected_components(mut points: HashSet<(i64, i64, i64)>, neighbors: &[(i64, i64, i64)]) -> Vec<HashSet<(i64, i64, i64)>> {
    let mut result = Vec::new();
    while let Some(&start) = points.iter().next() {
        result.push(connected_component_of(&mut points, start, neighbors));
    }
    result
}

fn main() {
    let input = std::fs::read_to_string("input.txt").ok().unwrap();

    // input points
    let points: Vec<(i64, i64, i64)> = input
        .split('\n')
        .filter(|line| !line.is_empty())
        .map(|line| {
            line
                .split(',')
                .map(|n| n.parse::<i64>().unwrap())
                .next_tuple()
                .unwrap()
        })
        .collect::<Vec<_>>();

    let points_set = HashSet::from_iter(points);

    let part_1 = points_set
        .iter()
        .flat_map(|&p| displace(p, &N6))
        .filter(|p| !points_set.contains(p))
        .count();

    println!("{}", part_1);

    let components = connected_components(points_set, &N14);

    // for each connected component of the input points, the shell is defined
    // as the set of air cubes that surround it, but are not enclosed in it.
    // This is equivalent to the connected component of the topmost air cell.
    let shells = components
        .iter()
        .map(|c| {
            let mut shell = c
                .iter()
                .flat_map(|&p| displace(p, &N14))
                .filter(|p| !c.contains(p))
                .collect::<HashSet<_>>();
            let max = *shell.iter().max().unwrap();
            connected_component_of(&mut shell, max, &N6)
        }).collect::<Vec<_>>();

    // the total number of faces of each component that interface with the
    // respective shell
    let counts = components
        .iter()
        .zip(shells.iter())
        .map(|(component, shell)| {
            component
                .iter()
                .flat_map(|&p| displace(p, &N6))
                .filter(|p| shell.contains(p))
                .count() as i64
        }).collect::<Vec<_>>();

    // for each cell of the input, the direction is either -1, 0 or 1.
    // -1 if it interfaces below (but not above) with its shell.
    //  1 if it interfaces above (but not below) with its shell.
    //  0 otherwise.
    //  This is equivalent to the number of shells that you exit when you cross
    //  this cell from below to above (negative if instead you are entering a shell).
    let dirs = components
        .iter()
        .zip(shells.iter())
        .flat_map(|(component, shell)| {
            component
                .iter()
                .map(|&p| (p, displace(p, &[(0, 0, -1), (0, 0, 1)]).tuple_windows().next().unwrap()))
                .map(|(p, (below, above))| (p, shell.contains(&above) as i64 - shell.contains(&below) as i64))
        }).sorted()
        .collect::<Vec<_>>();

    // the nesting depth of the connected component
    let depth = components
        .iter()
        .map(|component| {
            let &(x, y, z) = component.iter().next().unwrap();
            // the range of filled voxels that have the same x, y as the representative,
            // but a greater z coordinate.
            let lo = match dirs.binary_search(&((x, y, z), 0)) { Ok(n) => n + 1, Err(n) => n };
            let hi = match dirs.binary_search(&((x, y, i64::MAX), 0)) { Ok(n) => n, Err(n) => n };
            // the depth is given as the number of shells that we exit while traveling
            // towards positive z. The shell of the same connected component is not counted.
            dirs[lo..hi]
                .iter()
                .map(|(p, s)| if component.contains(p) {0} else {*s})
                .sum::<i64>()
        }).collect::<Vec<_>>();

    // the result for part 2 is the number of faces that interface connected components
    // of depth 0 with their respective shells.
    let part_2 = counts
        .iter()
        .zip(depth)
        .map(|(c, d)| if d == 0 {*c} else {0})
        .sum::<i64>();

    println!("{}", part_2);
}