Aoc_2025_Day8_Part2.rs

/// Advent of Code: Year 2025 Day 8 Part 2 Solver
/// ---------------------------------------------
/// 
/// Algorithm Idea
/// 
/// It works in 2 phases:
///   1. It "sort" all pairs of points by ascending distance
///   2. It uses the Union-Find algorithm to connect the circuits
///      The Union-Find algorithm is sometimes called Disjoint-Set Union
/// 
/// The "sorting" does not sort the entire list of every pairs.
/// Instead it "cuts" the 3D space into n^3 cubes called chunks
/// where n is the value passed to the command line argument "--chunks". 
/// Each axis X, Y and Z is cut into n pieces.
/// 
/// It helps finding rapidly points that are closed to each other because
/// they are in chunks that are aso close to each other.

use clap::Parser;
use std::cmp::Ordering;
use std::collections::{BinaryHeap, HashMap};
use std::io::prelude::*;
use std::time::Instant;
use std::{error::Error, fs::File, process::ExitCode};

const MAX_COORD: usize = 100000;

#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
struct Args {
    /// input file
    #[arg(short, long)]
    input: String,

    /// chunks
    #[arg(short, long, default_value_t = 5)] // Seem to be optimal with 5
    chunks: u32,
}

#[derive(Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Debug, Hash)]
struct Point {
    x: u32,
    y: u32,
    z: u32,
}

impl Point {
    /// Computing the square of the distance is enough and avoid floating
    /// point computations
    pub fn distance2(&self, other: &Self) -> i64 {
        let dx = (self.x as i64) - (other.x as i64);
        let dy = (self.y as i64) - (other.y as i64);
        let dz = (self.z as i64) - (other.z as i64);
        dx * dx + dy * dy + dz * dz
    }
}

/// A Pair of Points
#[derive(Clone, Copy, Debug)]
struct Pair {
    pub fst: Point,
    pub snd: Point,
    distance: i64,
}

impl Pair {
    #[inline(always)]
    fn new(fst: Point, snd: Point) -> Pair {
        Pair {
            fst,
            snd,
            distance: fst.distance2(&snd),
        }
    }
}

impl Eq for Pair {}

impl PartialEq for Pair {
    fn eq(&self, other: &Self) -> bool {
        self.fst == other.fst && self.snd == other.snd
    }
}

/// To find the pair of points with the lowest distance,
/// pairs will be stored in a BinaryHeap.
/// Unfortunately, this structure peeks the greatest item
/// but we want the lowest one so we invert the ordering
/// to get the minimal distance as the "greater" value
impl Ord for Pair {
    fn cmp(&self, other: &Self) -> Ordering {
        other.distance.cmp(&self.distance).then_with(|| {
            self.fst
                .cmp(&other.fst)
                .then_with(|| self.snd.cmp(&other.snd))
        })
    }
}

impl PartialOrd for Pair {
    #[inline(always)]
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

/// A chunk is a cubical region of space
#[repr(transparent)]
struct Chunk {
    /// The points in that region
    points: Vec<Point>,
}

/// Coordinate of a Chunk
/// 
/// This is not the same scale as points!
/// It is coordinates in the matrix of chunks.
#[derive(Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Debug, Hash)]
struct ChunkIndex {
    x: usize,
    y: usize,
    z: usize,
}

/// The pair of chunks needs to store the distance range of
/// all distances between points in the first and the second chunk of this pair
#[derive(Clone, Copy, Debug)]
struct ChunkPair {
    fst: ChunkIndex,
    snd: ChunkIndex,
    min_distance: i64,
    max_distance: i64,
}

impl Eq for ChunkPair {}

impl PartialEq for ChunkPair {
    fn eq(&self, other: &Self) -> bool {
        self.fst == other.fst && self.snd == other.snd
    }
}

/// Unlike Pair, we keep the natural ordering on distances because
/// Vec::sort sort in ascending order and we want lowest distances first.
impl Ord for ChunkPair {
    fn cmp(&self, other: &Self) -> Ordering {
        self.min_distance.cmp(&other.min_distance).then_with(|| {
            self.max_distance.cmp(&other.max_distance).then_with(|| {
                self.fst
                    .cmp(&other.fst)
                    .then_with(|| self.snd.cmp(&other.snd))
            })
        })
    }
}

impl PartialOrd for ChunkPair {
    #[inline(always)]
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

/// Iterate of all pairs of points
/// from the lowest distance to the greatest
/// 
/// It is semi-lazy:
/// "small" distances are loaded into the heap low_heap
/// When this heap is empty, the ordered list of chunk pairs is used
/// to fill it with the next "small" pairs.
/// 
/// The chunks don't align perfectly with distance separation.
/// But we need to have a clear distinction between "low" and "hight" distance
/// 
/// It thanks to this invariant:
/// 
/// If current_position < chunk_pairs.len()
/// then all distances in low_heap are <  to chunk_pairs[j].min_distance
/// and  all distances in high_vec are >= to chunk_pairs[j].min_distance
/// for all j >= current_position
/// 
/// It guarantees that the pairs in low_heap are the shortest.

struct PairIterator {
    nb_chunks: usize,
    nb_chunks_square: usize,

    chunks: Vec<Chunk>,

    nb_chunk_pairs: usize,
    chunk_pairs: Vec<ChunkPair>,

    current_position: usize,
    low_heap: BinaryHeap<Pair>,
    high_vec: Vec<Pair>
}

impl Iterator for PairIterator {
    type Item = Pair;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            // the low_heap contains the smallest pairs
            let r = self.low_heap.pop();

            if r.is_some() {
                return r;
            }

            // if true, the iterator has finished
            if self.current_position >= self.nb_chunk_pairs {
                return None;
            }

            // Refilling low_heap with previsouly distant pairs.
            for p in &self.high_vec {
                self.low_heap.push(*p);
            }
            self.high_vec = Vec::with_capacity(100000);

            // The next distance separation between low and high
            let up_to_distance = self.chunk_pairs[self.current_position].max_distance;

            // Filling all pairs whose distance is <= up_to_distance
            while self.current_position < self.nb_chunk_pairs
                && self.chunk_pairs[self.current_position].min_distance <= up_to_distance
            {
                let f = self.chunk_pairs[self.current_position].fst;
                let s = self.chunk_pairs[self.current_position].snd;

                // They have to be different
                for fst in
                    &self.chunks[f.x * self.nb_chunks_square + f.y * self.nb_chunks + f.z].points
                {
                    for snd in &self.chunks
                        [s.x * self.nb_chunks_square + s.y * self.nb_chunks + s.z]
                        .points
                    {
                        let pair = Pair::new(*fst, *snd);
                        // Remember that chunk contains pairs whose distance are > up_to_distance
                        // We need to consume them, but keep them apart from low_heap.
                        if pair.distance <= up_to_distance {
                            self.low_heap.push(pair);
                        } else {
                            self.high_vec.push(pair);
                        }
                    }
                }

                self.current_position += 1;
            }
        }
    }
}

struct UnionFind {
    union_find: HashMap<Point, Point>,
    roots: u32
}

impl UnionFind {
    pub fn new(_roots: u32) -> UnionFind {
        UnionFind {
            union_find: HashMap::with_capacity(1000),
            roots: _roots
        }
    }

    pub fn update(&mut self, k: Point, v: Point) {
        if !self.union_find.contains_key(&k) {
            self.roots -= 1;
        }
        self.union_find.insert(k, v);
    }

    fn find_rep(&mut self, k: Point) -> Point {
        match self.union_find.get_mut(&k) {
            Some(p) => {
                let p_owned = *p;
                let rep = self.find_rep(p_owned);
                self.update(k, rep);
                rep
            }
            None => k,
        }
    }

    pub fn connect(&mut self, p: &Pair) {
        let fst_rep = self.find_rep(p.fst);
        let snd_rep = self.find_rep(p.snd);
        if fst_rep != snd_rep {
            self.update(snd_rep, fst_rep);
        }
    }
}

#[inline(always)]
pub fn distance_range(fst: usize, snd: usize, width: usize) -> (i64, i64) {
    let w = width as i64;
    let (f, s) = {
        if fst <= snd {
            (fst as i64, snd as i64)
        } else {
            (snd as i64, fst as i64)
        }
    };
    (i64::max((s - f) * w - 1, 0), (s + 1 - f) * w + -1)
}

fn main() -> Result<ExitCode, Box<dyn Error>> {
    let args = Args::parse();

    let mut file = File::open(&args.input)?;
    let mut content: String = String::new();
    let _ = file.read_to_string(&mut content)?;

    let points = {
        let mut pts = Vec::new();

        for line in content.split("\n") {
            let trimmed = line.trim();

            if trimmed.is_empty() {
                continue;
            }

            let t: Vec<&str> = trimmed.split(",").collect();
            pts.push(Point {
                x: t[0].parse::<u32>()?,
                y: t[1].parse::<u32>()?,
                z: t[2].parse::<u32>()?,
            });
        }

        pts
    };

    // Input Parsed
    let start = Instant::now();

    let nb_chunks: usize = args.chunks as usize;
    let nb_chunks_square: usize = nb_chunks * nb_chunks;
    let nb_chunks_cube: usize = nb_chunks_square * nb_chunks;

    let chunk_step: usize = MAX_COORD / nb_chunks;

    let mut chunks: Vec<Chunk> = Vec::with_capacity(nb_chunks_cube);

    // Initialize chunks
    for _i in 0..nb_chunks_cube {
        chunks.push(Chunk { points: Vec::new() });
    }

    let mut nb_points: u32 = 0;

    // Classify points in chunks
    for p in points {
        nb_points += 1;
        chunks[((p.x as usize) / chunk_step) * nb_chunks_square
            + ((p.y as usize) / chunk_step) * nb_chunks
            + ((p.z as usize) / chunk_step)]
            .points
            .push(p);
    }

    // Distance 0 to chunk_step
    let mut low_heap = BinaryHeap::new();
    let mut high_vec = Vec::new();
    let up_to_distance = (chunk_step as i64) * (chunk_step as i64);

    // Same chunk
    for x in 0..nb_chunks {
        for y in 0..nb_chunks {
            for z in 0..nb_chunks {
                let pts = &chunks[x * nb_chunks_square + y * nb_chunks + z].points;
                let n = pts.len();

                for i in 0..n {
                    for j in i + 1..n {
                        let pair = Pair::new(pts[i], pts[j]);
                        if pair.distance <= up_to_distance {
                            low_heap.push(pair);
                        } else {
                            high_vec.push(pair);
                        }
                    }
                }
            }
        }
    }

    // Adjacent Chunks
    for x1 in 0..nb_chunks {
        for y1 in 0..nb_chunks {
            for z1 in 0..nb_chunks {
                let fst_pts = &chunks[x1 * nb_chunks_square + y1 * nb_chunks + z1].points;

                for x2 in x1..usize::min(x1 + 2, nb_chunks) {
                    for y2 in y1..usize::min(y1 + 2, nb_chunks) {
                        for z2 in z1..usize::min(z1 + 2, nb_chunks) {
                            if x1 != x2 || y1 != y2 || z1 != z2 {
                                for fst in fst_pts {
                                    for snd in
                                        &chunks[x2 * nb_chunks_square + y2 * nb_chunks + z2].points
                                    {
                                        let pair = Pair::new(*fst, *snd);
                                        if pair.distance <= up_to_distance {
                                            low_heap.push(pair);
                                        } else {
                                            high_vec.push(pair);
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    // Other chunk pairs
    let mut chunk_pairs = Vec::new();

    for x1 in 0..nb_chunks {
        for y1 in 0..nb_chunks {
            for z1 in 0..nb_chunks {
                let fst = ChunkIndex {
                    x: x1,
                    y: y1,
                    z: z1,
                };

                for x2 in x1..nb_chunks {
                    for y2 in y1..nb_chunks {
                        for z2 in z1..nb_chunks {
                            if x2 - x1 > 1 || y2 - y1 > 1 || z2 - z1 > 1 {
                                let snd = ChunkIndex {
                                    x: x2,
                                    y: y2,
                                    z: z2,
                                };

                                let (dx_min, dx_max) = distance_range(x1, x2, chunk_step);
                                let (dy_min, dy_max) = distance_range(y1, y2, chunk_step);
                                let (dz_min, dz_max) = distance_range(z1, z2, chunk_step);

                                let pair = ChunkPair {
                                    fst,
                                    snd,
                                    min_distance: dx_min * dx_min
                                        + dy_min * dy_min
                                        + dz_min * dz_min,
                                    max_distance: dx_max * dx_max
                                        + dy_max * dy_max
                                        + dz_max * dz_max,
                                };

                                chunk_pairs.push(pair)
                            }
                        }
                    }
                }
            }
        }
    }

    // Minimal distances first
    chunk_pairs.sort();

    // Let's iterate over pairs in distance ascending order
    let mut pair_iterator = PairIterator {
        nb_chunks,
        nb_chunks_square,
        chunks,
        nb_chunk_pairs: chunk_pairs.len(),
        chunk_pairs,
        current_position: 0,
        low_heap: low_heap,
        high_vec: high_vec
    };

    let mut union_find = UnionFind::new(nb_points);
    let mut last_pair: Option<Pair> = None;

    while let Some(p) = &pair_iterator.next() {
        if union_find.roots == 1 {
            break;
        }

        union_find.connect(&p);
        last_pair = Some(*p);
    }

    let answer = match last_pair {
        Some(p) => p.fst.x * p.snd.x,
        None => panic!("No last pair"),
    };

    let stop = Instant::now();

    println!(
        "Part 2: {}\nTotal time: {} micros",
        answer,
        stop.duration_since(start).as_micros()
    );

    Ok(ExitCode::SUCCESS)
}