Stick closer to algorithm 2 from the paper

This commit is contained in:
Dirkjan Ochtman 2021-01-07 21:19:50 +01:00
parent f01ed4a4a0
commit 6def318423
1 changed files with 14 additions and 20 deletions

View File

@ -430,8 +430,10 @@ trait Layer {
/// case, we use `links` to constrain the number of per-candidate links we consider for search. /// case, we use `links` to constrain the number of per-candidate links we consider for search.
fn search<P: Point>(&self, point: &P, search: &mut Search, points: &[P], links: usize) { fn search<P: Point>(&self, point: &P, search: &mut Search, points: &[P], links: usize) {
while let Some(Reverse(candidate)) = search.candidates.pop() { while let Some(Reverse(candidate)) = search.candidates.pop() {
if candidate.distance > search.furthest { if let Some(furthest) = search.nearest.last() {
break; if candidate.distance > furthest.distance {
break;
}
} }
for pid in self.nearest_iter(candidate.pid).take(links) { for pid in self.nearest_iter(candidate.pid).take(links) {
@ -459,8 +461,6 @@ pub struct Search {
nearest: Vec<Candidate>, nearest: Vec<Candidate>,
/// Maximum number of nearest neighbors to retain (`ef` in the paper) /// Maximum number of nearest neighbors to retain (`ef` in the paper)
ef: usize, ef: usize,
/// Current furthest node in `nearest`
furthest: OrderedFloat<f32>,
} }
impl Search { impl Search {
@ -470,14 +470,12 @@ impl Search {
visited, visited,
candidates, candidates,
nearest, nearest,
furthest,
ef: _, ef: _,
} = self; } = self;
visited.clear(); visited.clear();
candidates.clear(); candidates.clear();
nearest.clear(); nearest.clear();
*furthest = OrderedFloat::from(f32::INFINITY);
} }
/// Track node `pid` as a potential new neighbor for the given `point` /// Track node `pid` as a potential new neighbor for the given `point`
@ -491,20 +489,18 @@ impl Search {
let other = &points[pid]; let other = &points[pid];
let distance = OrderedFloat::from(point.distance(other)); let distance = OrderedFloat::from(point.distance(other));
if self.nearest.len() >= self.ef && distance > self.furthest {
return;
}
if self.nearest.len() > self.ef * 2 {
self.nearest.sort_unstable();
self.nearest.truncate(self.ef);
self.furthest = self.nearest.last().unwrap().distance;
}
let new = Candidate { distance, pid }; let new = Candidate { distance, pid };
let idx = match self.nearest.binary_search(&new) {
Err(idx) if idx < self.ef => idx,
Err(_) => return,
Ok(_) => unreachable!(),
};
self.nearest.insert(idx, new);
self.candidates.push(Reverse(new)); self.candidates.push(Reverse(new));
self.nearest.push(new); if self.nearest.len() > self.ef {
self.furthest = max(self.furthest, distance); self.nearest.truncate(self.ef);
}
} }
/// Lower the search to the next lower level /// Lower the search to the next lower level
@ -516,7 +512,6 @@ impl Search {
/// because `Layer::search()` is always called right before calling `cull()`. /// because `Layer::search()` is always called right before calling `cull()`.
fn cull(&mut self) { fn cull(&mut self) {
self.nearest.truncate(self.ef); // Limit size of the set of nearest neighbors self.nearest.truncate(self.ef); // Limit size of the set of nearest neighbors
self.furthest = self.nearest.last().unwrap().distance;
self.candidates.clear(); self.candidates.clear();
for &candidate in self.nearest.iter() { for &candidate in self.nearest.iter() {
@ -535,7 +530,6 @@ impl Default for Search {
candidates: BinaryHeap::new(), candidates: BinaryHeap::new(),
nearest: Vec::new(), nearest: Vec::new(),
ef: 1, ef: 1,
furthest: OrderedFloat::from(f32::INFINITY),
} }
} }
} }