aoc-core 0.1.2

Useful Advent of Code data structures, types and functions common to my Rust solutions.
Documentation
/// Disjoint-set union-find data structure.
pub struct Dsu {
    /// Parent pointer for each element.
    ///
    /// A root element is its own parent.
    parent: Vec<usize>,

    /// Size of each component (valid only for roots).
    size: Vec<usize>,

    /// Current number of disjoint components.
    components: usize,
}

impl Dsu {
    /// Creates a new DSU with `n` singleton components.
    #[must_use]
    pub fn new(n: usize) -> Self {
        Self {
            parent: (0..n).collect(),
            size: vec![1; n],
            components: n,
        }
    }

    /// Parent point for each element.
    ///
    /// A root element is its own parent.
    #[inline]
    #[must_use]
    pub fn parent(&self) -> &[usize] {
        &self.parent
    }

    /// Size of each component (valid only for roots).
    #[inline]
    #[must_use]
    pub fn size(&self) -> &[usize] {
        &self.size
    }

    /// Current number of disjoint components.
    #[inline]
    #[must_use]
    pub const fn components(&self) -> usize {
        self.components
    }

    /// Returns the root of the component containing `x`.
    ///
    /// Applies path compression, making future queries faster.
    #[inline]
    pub fn find(&mut self, mut x: usize) -> usize {
        let mut root = x;

        while self.parent[root] != root {
            root = self.parent[root];
        }

        while self.parent[x] != x {
            let next = self.parent[x];
            self.parent[x] = root;
            x = next;
        }

        root
    }

    /// Unites the components containing `x` and `y`,
    /// merging the smaller component into the bigger one.
    ///
    /// Returns `true` if a merge occurred,
    /// or `false` if they were already in the same component.
    #[inline]
    pub fn union(&mut self, x: usize, y: usize) -> bool {
        let rx = self.find(x);
        let ry = self.find(y);

        if rx == ry {
            return false;
        }

        if self.size[rx] < self.size[ry] {
            self.parent[rx] = ry;
            self.size[ry] += self.size[rx];
        } else {
            self.parent[ry] = rx;
            self.size[rx] += self.size[ry];
        }

        self.components -= 1;

        true
    }
}

#[cfg(test)]
mod tests {
    use super::Dsu;

    #[test]
    fn new() {
        let input = 8;
        let expected = (vec![0, 1, 2, 3, 4, 5, 6, 7], vec![1; 8], 8);
        let dsu = Dsu::new(input);
        let output = (dsu.parent, dsu.size, dsu.components);
        assert_eq!(expected, output, "\n input: {input:?}");
    }

    // ------------------------------------------------------------------------------------------------
    // Find

    #[test]
    fn find() {
        let input = [0, 3, 1, 2];
        let expected = [0, 3, 0, 0];
        let mut dsu = Dsu {
            parent: vec![0, 0, 1, 3],
            size: vec![2, 1, 1, 1],
            components: 2,
        };
        let output = input.map(|n| dsu.find(n));
        assert_eq!(expected, output, "\n input: {input:?}");
    }

    #[test]
    fn find_compression() {
        let input = 2;
        let expected = 0;
        let mut dsu = Dsu {
            parent: vec![0, 0, 1, 3],
            size: vec![2, 1, 1, 1],
            components: 2,
        };
        let _ = dsu.find(input);
        let output = dsu.parent()[2];
        assert_eq!(expected, output, "\n input: {input:?}");
    }

    // ------------------------------------------------------------------------------------------------
    // Union

    #[test]
    fn union_size() {
        let input = [((0, 1), 0), ((2, 0), 0), ((3, 0), 0), ((3, 1), 0)];
        let expected = [(true, 2), (true, 3), (true, 4), (false, 4)];
        let mut dsu = Dsu::new(4);
        let output = input.map(|((x, y), r)| {
            let union = dsu.union(x, y);
            let root = dsu.find(r);
            (union, dsu.size()[root])
        });
        assert_eq!(expected, output, "\n input: {input:?}");
    }

    #[test]
    fn union_components() {
        let input = [(0, 1), (2, 3), (1, 3), (0, 2)];
        let expected = [(true, 3), (true, 2), (true, 1), (false, 1)];
        let mut dsu = Dsu::new(4);
        let output = input.map(|(x, y)| (dsu.union(x, y), dsu.components()));
        assert_eq!(expected, output, "\n input: {input:?}");
    }
}