#![no_std]
doc_comment::doctest!("../README.md");
#[cfg(feature = "alloc")]
extern crate alloc;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
#[cfg(feature = "alloc")]
use alloc::vec::Vec;
use num_traits::Zero;
pub trait Metric<P> {
type Unit: Ord + Zero + Copy;
fn distance(&self, a: &P, b: &P) -> Self::Unit;
}
#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct Neighbor<Unit, Ix = usize> {
pub index: Ix,
pub distance: Unit,
}
pub trait Knn<'a> {
type Ix: Copy;
type Point: 'a;
type Value: 'a;
type Metric: Metric<Self::Point>;
type KnnIter: IntoIterator<
Item = (
Neighbor<<Self::Metric as Metric<Self::Point>>::Unit, Self::Ix>,
&'a Self::Point,
&'a Self::Value,
),
>;
fn point(&self, index: Self::Ix) -> &'a Self::Point;
fn value(&self, index: Self::Ix) -> &'a Self::Value;
fn knn(&'a self, query: &Self::Point, num: usize) -> Self::KnnIter;
#[allow(clippy::type_complexity)]
fn nn(
&'a self,
query: &Self::Point,
) -> Option<(
Neighbor<<Self::Metric as Metric<Self::Point>>::Unit, Self::Ix>,
&'a Self::Point,
&'a Self::Value,
)>;
}
pub trait RangeQuery<'a>: Knn<'a> {
type RangeIter: IntoIterator<
Item = (
Neighbor<<Self::Metric as Metric<Self::Point>>::Unit, Self::Ix>,
&'a Self::Point,
&'a Self::Value,
),
>;
#[allow(clippy::type_complexity)]
fn range_query(
&self,
query: &Self::Point,
range: <Self::Metric as Metric<Self::Point>>::Unit,
) -> Self::RangeIter;
}
pub trait KnnInsert<'a>: Knn<'a> {
fn insert(&mut self, key: Self::Point, value: Self::Value) -> Self::Ix;
}
pub trait KnnFromMetricAndBatch<M, B> {
fn from_metric_and_batch(metric: M, batch: B) -> Self;
}
pub trait KnnFromBatch<M, B>: KnnFromMetricAndBatch<M, B> {
fn from_batch(batch: B) -> Self;
}
impl<M, B, T> KnnFromBatch<M, B> for T
where
T: KnnFromMetricAndBatch<M, B>,
M: Default,
{
fn from_batch(batch: B) -> Self {
Self::from_metric_and_batch(M::default(), batch)
}
}
#[cfg(feature = "alloc")]
pub struct LinearKnn<M, I> {
pub metric: M,
pub points: I,
}
#[cfg(feature = "alloc")]
impl<'a, M: Metric<P>, I, P: 'a, V: 'a> Knn<'a> for LinearKnn<M, I>
where
I: Iterator<Item = &'a (P, V)> + Clone,
{
type Ix = usize;
type Metric = M;
type Point = P;
type Value = V;
type KnnIter = Vec<(Neighbor<M::Unit>, &'a P, &'a V)>;
fn point(&self, index: Self::Ix) -> &'a Self::Point {
&self.points.clone().nth(index).unwrap().0
}
fn value(&self, index: Self::Ix) -> &'a Self::Value {
&self.points.clone().nth(index).unwrap().1
}
fn knn(&'a self, query: &Self::Point, num: usize) -> Self::KnnIter {
let mut dataset = self.points.clone().enumerate().map(|(index, (pt, val))| {
(
Neighbor {
index,
distance: self.metric.distance(pt, query),
},
pt,
val,
)
});
let mut neighbors = Vec::with_capacity(num);
neighbors.extend((&mut dataset).take(num));
neighbors.sort_unstable_by_key(|n| n.0.distance);
for point in dataset {
let position = neighbors.partition_point(|n| n.0.distance <= point.0.distance);
if position != num {
neighbors.pop();
neighbors.insert(position, point);
}
}
neighbors
}
#[allow(clippy::type_complexity)]
fn nn(
&self,
query: &Self::Point,
) -> Option<(
Neighbor<<Self::Metric as Metric<Self::Point>>::Unit, Self::Ix>,
&'a Self::Point,
&'a Self::Value,
)> {
self.points
.clone()
.enumerate()
.map(|(index, (pt, val))| {
(
Neighbor {
index,
distance: self.metric.distance(pt, query),
},
pt,
val,
)
})
.min_by_key(|n| n.0.distance)
}
}
#[cfg(feature = "alloc")]
impl<'a, M, I> KnnFromMetricAndBatch<M, I> for LinearKnn<M, I>
where
M: Default,
{
fn from_metric_and_batch(metric: M, points: I) -> Self {
Self { metric, points }
}
}