steiner-tree 0.0.5

// SPDX-FileCopyrightText: 2022 Thomas Kramer <code@tkramer.ch>
//
// SPDX-License-Identifier: GPL-3.0-or-later

//! Compute minimum spanning trees (MST) of points in a multi-dimensional space
//! with a defined distance metric.

mod locality_sensitive_hash;
mod nearest_neighbor;

use locality_sensitive_hash::*;
use nearest_neighbor::*;

use num_traits::Zero;
use std::cmp::Ordering;
use std::collections::{BinaryHeap, HashMap, HashSet};
use std::hash::Hash;

pub fn minimum_spanning_tree_hamming_space<P>(points: &[P], exact: bool) -> HashMap<P, P>
where
    P: HammingPoint<Coord = bool> + Clone + Hash + Eq,
{
    if points.is_empty() {
        return Default::default();
    }

    // Build locality sensitive hash.
    let lsh = {
        let bin_size = 8;
        let num_points = points.len() as f64;
        let num_bins = num_points / (bin_size as f64);
        let num_sub_samples = num_bins.log2().ceil() as usize;
        let lsh: HammingLSH<_, usize> = HammingLSH::new(num_sub_samples, &points[1..]);
        lsh
    };

    // Build nearest-neighbor search.
    let mut search = LSHNearestNeighbourSearch::new(lsh);

    for p in &points[1..] {
        search.insert(p.clone());
    }

    let mut front = BinaryHeap::new();
    front.push(PQElement {
        distance: 0u32,
        point: points[0].clone(),
        prev: points[0].clone(),
    });

    // Edges of the minimum spanning tree.
    let mut tree_edges = HashMap::new();

    // Store potential shortest edges.
    // If a point `p` is added to the tree, then all other nodes who had `n` as a nearest neighbour
    // need to get another nearest-neighbour candidate.
    let mut potential_prev_nodes: HashMap<P, Vec<P>> = HashMap::new();

    while let Some(e) = front.pop() {
        if tree_edges.contains_key(&e.point) {
            continue;
        }

        println!(
            "progress: {:.1} %",
            (100. * (tree_edges.len() as f64) / (points.len() as f64))
        );

        search.remove(&e.point);

        // All nodes who had `e.point` as nearest neighbour need to be assigned with new nearest neighbours.
        {
            let users_of_p = potential_prev_nodes.remove(&e.point).unwrap_or(Vec::new());
            // `e.point` was a nearest neighbour of itself. Hence it needs to be assigned with a yes unused nearest neighbour.
            let users_of_p = std::iter::once(e.point.clone()).chain(users_of_p); // Skip already used nodes.

            // Find yet unused nearest neighbours for all points who had `e.point` as nearest neighbour.
            for p in users_of_p {
                let maybe_nearest_neighbour = if exact {
                    search.exact_nearest_neighbour(&p)
                } else {
                    search
                        .approx_nearest_neighbour(&p)
                        .or_else(|| search.exact_nearest_neighbour(&p))
                };

                if let Some(n) = maybe_nearest_neighbour {
                    let dist = p.distance(n);

                    front.push(PQElement {
                        distance: dist,
                        point: n.clone(),
                        prev: p.clone(),
                    });

                    potential_prev_nodes
                        .entry(n.clone())
                        .or_insert(Vec::new())
                        .push(p);
                }
            }
        }

        println!("dist = {}", e.point.distance(&e.prev));
        // Create edge.
        if e.point != e.prev {
            // Skip the edge from the first point to itself.
            tree_edges.insert(e.point, e.prev);
        }
    }

    tree_edges
}

/// Element in the priority queue.
/// Used to sort edges by ascending distance.
struct PQElement<D, V> {
    distance: D,
    point: V,
    prev: V,
}

impl<D: Ord + Eq, V> Ord for PQElement<D, V> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.distance.cmp(&other.distance).reverse()
    }
}

impl<D: PartialOrd, V> PartialOrd for PQElement<D, V> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.distance
            .partial_cmp(&other.distance)
            .map(|o| o.reverse())
    }
}

impl<D: PartialEq, V> Eq for PQElement<D, V> {}

impl<D: PartialEq, V> PartialEq for PQElement<D, V> {
    fn eq(&self, other: &Self) -> bool {
        self.distance.eq(&other.distance)
    }
}

#[test]
fn test_minimum_spanning_tree_with_hamming_distance_of_integers() {
    // Under Hamming distance metric the minimum spanning tree of natural numbers (with zero)
    // until `n` is equal to a 'Gray code'. E.g. a sequence of integers where neigbors differ only
    // by one bit.
    let num_points = 100;
    let points: Vec<_> = (0..num_points).collect();

    let tree_edges = minimum_spanning_tree_hamming_space(&points, true);

    dbg!(&tree_edges);

    for (a, b) in tree_edges.iter() {
        let dist = a.distance(b);
        assert_eq!(dist, 1);
    }

    assert_eq!(tree_edges.len(), num_points - 1);
}