Skip to main content

knn

Function knn 

Source
pub fn knn<T: VectorType, M: DistanceMetric<T>>(
    data: impl Indexable<T>,
    query: &[T],
    k: usize,
) -> Vec<(f32, usize)>
Expand description

Performs a K-Nearest Neighbors (KNN) search in parallel.

This function finds the k closest vectors in the data to the given query vector, using the specified distance metric M. It utilizes data parallelization to efficiently compare entries and keeps track of the nearest neighbors.

§Arguments

  • data - The dataset to search against. Must implement Indexable<T>. Common usage includes passing a tuple of (&[T], dim) for flattened arrays, or passing nested vectors.
  • query - The target query vector.
  • k - The maximum number of nearest neighbors to retrieve.

§Returns

A Vec<(f32, usize)> sorted from the closest to the furthest distance. Each element contains the calculated distance and the index of the neighbor.

§Example

use flat_knn::{knn, L2};

let data = vec![
    1.0, 2.0, 3.0, 4.0, // index 0 (dist = 1)
    8.0, 7.0, 6.0, 5.0, // index 1 (dist = 84)
    1.0, 2.0, 3.0, 9.0, // index 2 (dist = 16)
];
let dim = 4;
let query = [1.0, 2.0, 3.0, 5.0];

// Find 2 nearest neighbors using L2 distance
let neighbors = knn::<_, L2>((&data, dim), &query, 2);

assert_eq!(neighbors.len(), 2);
assert_eq!(neighbors[0], (1.0, 0));
assert_eq!(neighbors[1], (16.0, 2));