diff --git a/src/float/query/nearest_n_within.rs b/src/float/query/nearest_n_within.rs index 3fc1d06..f8fe7bf 100644 --- a/src/float/query/nearest_n_within.rs +++ b/src/float/query/nearest_n_within.rs @@ -230,6 +230,71 @@ mod tests { } } + #[test] + fn can_query_nearest_n_items_unsorted_max_qty() { + let mut tree: KdTree = KdTree::new(); + + let content_to_add: [([AX; 4], u32); 16] = [ + ([0.9f32, 0.0f32, 0.9f32, 0.0f32], 9), + ([0.4f32, 0.5f32, 0.4f32, 0.5f32], 4), + ([0.12f32, 0.3f32, 0.12f32, 0.3f32], 12), + ([0.7f32, 0.2f32, 0.7f32, 0.2f32], 7), + ([0.13f32, 0.4f32, 0.13f32, 0.4f32], 13), + ([0.6f32, 0.3f32, 0.6f32, 0.3f32], 6), + ([0.2f32, 0.7f32, 0.2f32, 0.7f32], 2), + ([0.14f32, 0.5f32, 0.14f32, 0.5f32], 14), + ([0.3f32, 0.6f32, 0.3f32, 0.6f32], 3), + ([0.10f32, 0.1f32, 0.10f32, 0.1f32], 10), + ([0.16f32, 0.7f32, 0.16f32, 0.7f32], 16), + ([0.1f32, 0.8f32, 0.1f32, 0.8f32], 1), + ([0.15f32, 0.6f32, 0.15f32, 0.6f32], 15), + ([0.5f32, 0.4f32, 0.5f32, 0.4f32], 5), + ([0.8f32, 0.1f32, 0.8f32, 0.1f32], 8), + ([0.11f32, 0.2f32, 0.11f32, 0.2f32], 11), + ]; + + for (point, item) in content_to_add { + tree.add(&point, item); + } + + let query_point = [0.78f32, 0.55f32, 0.78f32, 0.55f32]; + + let radius = 100.0; + let max_qty = 1; + + let result_unsorted: Vec<_> = tree + .nearest_n_within::(&query_point, radius, max_qty, false) + .into_iter() + .map(|n| (n.distance, n.item)) + .collect(); + + let result_sorted: Vec<_> = tree + .nearest_n_within::(&query_point, radius, max_qty, true) + .into_iter() + .map(|n| (n.distance, n.item)) + .collect(); + + assert_eq!(result_unsorted.len(), max_qty); + assert_eq!(result_sorted.len(), max_qty); + + let max_qty = 0; + + let result_unsorted: Vec<_> = tree + .nearest_n_within::(&query_point, radius, max_qty, false) + .into_iter() + .map(|n| (n.distance, n.item)) + .collect(); + + let result_sorted: Vec<_> = tree + .nearest_n_within::(&query_point, radius, max_qty, true) + .into_iter() + .map(|n| (n.distance, n.item)) + .collect(); + + assert_eq!(result_unsorted.len(), max_qty); + assert_eq!(result_sorted.len(), max_qty); + } + #[test] fn can_query_nearest_n_items_within_radius_unsorted_large_scale() { const TREE_SIZE: usize = 100_000;