steiner-tree 0.0.5

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

//! Nearest neighbor search based on locality sensitive hashing (LSH).
//! TODO: Use multi-index hash tables.

use itertools::Itertools;

use std::collections::HashMap;
use std::hash::Hash;

use super::locality_sensitive_hash::*;

/// Accelerate nearest-neighbor search by putting close points into the same bin.
struct HashBins<T, LSH, H> {
    bins: HashMap<H, Vec<T>>,
    hasher: LSH,
}

impl<T, LSH, H> HashBins<T, LSH, H>
where
    H: Hash + Eq,
    LSH: LocalitySensitiveHasher<Domain = T, Image = H>,
{
    pub fn new(hasher: LSH) -> Self {
        Self {
            bins: Default::default(),
            hasher,
        }
    }

    pub fn insert(&mut self, value: T) {
        let hash = self.hasher.hash(&value);
        self.bins.entry(hash).or_default().push(value)
    }

    /// Iterate over the bin which contains the given `value`.
    pub fn bin_content(&self, value: &T) -> impl Iterator<Item = &T> {
        let hash = self.hasher.hash(value);
        self.bins.get(&hash).into_iter().flat_map(|bin| bin.iter())
    }

    /// Get bins sorted by distance to `value`.
    fn bins_sorted_by_distance(
        &self,
        value: &T,
    ) -> impl Iterator<Item = (LSH::ImageDistance, &Vec<T>)> {
        let hash = self.hasher.hash(value);
        let mut other_bins: Vec<_> = self
            .bins
            .iter()
            .map(|(bin_hash, bin)| {
                let dist = self.hasher.image_distance(bin_hash, &hash);
                (dist, bin)
            })
            .collect();

        /// Sort bins by distance to `hash`.
        other_bins.sort_by(|(h1, _), (h2, _)| h1.cmp(h2));

        other_bins.into_iter()
    }

    pub fn bin_contents_sorted_by_distance(
        &self,
        value: &T,
    ) -> impl Iterator<Item = (LSH::ImageDistance, impl Iterator<Item = &T>)> {
        self.bins_sorted_by_distance(value)
            .map(|(dist, bin_content)| (dist, bin_content.iter()))
    }
}

impl<T, LSH, H> HashBins<T, LSH, H>
where
    H: Hash + Eq,
    LSH: LocalitySensitiveHasher<Domain = T, Image = H>,
    T: PartialEq,
{
    /// Remove an element from its bin.
    pub fn remove(&mut self, value: &T) -> Option<T> {
        let hash = self.hasher.hash(&value);

        let removed = self.bins.get_mut(&hash).and_then(|bin| {
            if let Some((idx, _)) = bin.iter().find_position(|x| *x == value) {
                Some(bin.swap_remove(idx))
            } else {
                None
            }
        });

        if let Some(bin) = self.bins.get(&hash) {
            if bin.is_empty() {
                self.bins.remove(&hash);
            }
        }

        removed
    }
}

// pub trait NearestNeighbourSearch {
//     type Point;
//
//     fn nearest_neighbours(&self, p: &Self::Point) -> impl Iterator<Item=&Self::Point>;
// }

trait Distance {
    type D: Ord;
}

/// Nearest-neighbour search using locality sensitive hashing (LSH) for acceleration.
pub struct LSHNearestNeighbourSearch<T, LSH, H> {
    hash_bins: HashBins<T, LSH, H>,
}

impl<T, LSH, H> LSHNearestNeighbourSearch<T, LSH, H>
where
    H: Hash + Eq,
    LSH: LocalitySensitiveHasher<Domain = T, Image = H>,
{
    pub fn new(hasher: LSH) -> Self {
        Self {
            hash_bins: HashBins::new(hasher),
        }
    }

    pub fn insert(&mut self, value: T) {
        self.hash_bins.insert(value)
    }
}

impl<T, LSH, H> LSHNearestNeighbourSearch<T, LSH, H>
where
    T: PartialEq,
    H: Hash + Eq,
    LSH: LocalitySensitiveHasher<Domain = T, Image = H>,
{
    pub fn remove(&mut self, value: &T) -> Option<T> {
        self.hash_bins.remove(value)
    }
}

impl<T, LSH, H, Dist> LSHNearestNeighbourSearch<T, LSH, H>
where
    H: Hash + Eq,
    LSH:
        LocalitySensitiveHasher<Domain = T, Image = H, DomainDistance = Dist, ImageDistance = Dist>,
    T: PartialEq,
    Dist: Ord,
{
    /// Find nearest neighbour among the points which are in the same
    /// bin as the `value`, if any.
    pub fn approx_nearest_neighbour(&self, value: &T) -> Option<&T> {
        self.hash_bins
            .bin_content(value)
            .filter(|&x| x != value)
            // Find element with minimal distance to x.
            .min_by_key(|x| self.hash_bins.hasher.domain_distance(value, x))
    }

    /// Find nearest neighbour among all points.
    /// The search is started from the closest bin.
    /// The search can be terminated once the minimum distance from a bin to the `value` exceeds
    /// the distance to the current closest neighbour.
    pub fn exact_nearest_neighbour(&self, value: &T) -> Option<&T> {
        let neigbhour_bins = self.hash_bins.bin_contents_sorted_by_distance(value);

        let mut nearest_neighbour = None;
        let mut nearest_neighbour_distance = None;

        for (bin_distance, bin_content) in neigbhour_bins {
            // Terminate search when no better result can be found.
            if let Some(best_distance) = &nearest_neighbour_distance {
                if &bin_distance >= best_distance {
                    break;
                }
            }

            let nearest_in_bin = bin_content
                .filter(|&x| x != value)
                // Find element with minimal distance to x.
                .min_by_key(|x| self.hash_bins.hasher.domain_distance(value, x));

            if let Some(n) = nearest_in_bin {
                let dist = self.hash_bins.hasher.domain_distance(value, n);

                let is_better = nearest_neighbour_distance
                    .as_ref()
                    .map(|d| &dist < d)
                    .unwrap_or(true);

                if is_better {
                    nearest_neighbour = Some(n);
                    nearest_neighbour_distance = Some(dist);
                }
            }
        }

        nearest_neighbour
    }
}

#[test]
fn test_nearest_neighbour_search() {
    let points = [0b0000, 0b0001, 0b0010, 0b0011, 0b0100, 0b0110, 0b0111];

    let lsh: HammingLSH<_, usize> = HammingLSH::new(1, &points);

    let mut search = LSHNearestNeighbourSearch::new(lsh);

    for p in points {
        search.insert(p);
    }

    assert_eq!(search.exact_nearest_neighbour(&0b0000), Some(&0b0001));
    assert_eq!(search.exact_nearest_neighbour(&0b0111), Some(&0b0110));
}