use crate::ID;
#[derive(Debug)]
pub struct SearchResult {
pub similarity: f32,
pub id: ID,
}
#[derive(Debug)]
pub struct ResultSet {
sims: Vec<f32>,
ids: Vec<ID>,
k: usize,
pub checked: usize,
}
impl ResultSet {
pub fn new(k: usize) -> Self {
Self {
sims: Vec::with_capacity(k),
ids: Vec::with_capacity(k),
k,
checked: 0,
}
}
pub fn len(&self) -> usize {
self.sims.len()
}
pub fn is_empty(&self) -> bool {
self.sims.is_empty()
}
pub fn compute_recall(&self, baseline: &ResultSet, at: usize) -> f64 {
let mut found = 0;
for x in baseline.ids.iter().take(at) {
for y in self.ids.iter().take(at) {
if x == y {
found += 1;
}
}
}
found as f64 / at as f64
}
pub fn add_result(&mut self, id: ID, similarity: f32) {
self.checked += 1;
if self.sims.len() == self.k {
let last = self.sims.last().unwrap_or(&f32::MIN);
if *last > similarity {
return;
}
}
let mut insert: usize = 0;
let mut found: bool = false;
for (self_id, self_sim) in self.ids.iter().zip(self.sims.iter()) {
if id == *self_id {
return;
}
if *self_sim < similarity {
found = true;
break;
}
insert += 1;
}
if !found {
if self.sims.len() < self.k {
self.sims.push(similarity);
self.ids.push(id);
}
return;
}
self.ids.insert(insert, id);
self.ids.truncate(self.k);
self.sims.insert(insert, similarity);
self.sims.truncate(self.k);
}
pub fn iter_results(&self) -> impl Iterator<Item = SearchResult> + '_ {
self.sims
.iter()
.zip(self.ids.iter())
.map(|(sim, id)| SearchResult {
similarity: *sim,
id: *id,
})
}
}