use crate::algorithms::spatial::edge_lengths::DistanceMetric;
use crate::core::{Graph, IgraphError, IgraphResult};
#[allow(clippy::many_single_char_names)]
pub fn nearest_neighbor_graph(
points: &[Vec<f64>],
metric: DistanceMetric,
k: i64,
cutoff: f64,
directed: bool,
) -> IgraphResult<Graph> {
let n = points.len();
if n == 0 {
return Graph::new(0, directed);
}
let dim = points[0].len();
if dim == 0 {
return Err(IgraphError::InvalidArgument(
"nearest_neighbor_graph: 0-dimensional points are not supported".into(),
));
}
for (i, row) in points.iter().enumerate().skip(1) {
if row.len() != dim {
return Err(IgraphError::InvalidArgument(format!(
"nearest_neighbor_graph: point row {i} has dimension {} but expected {dim}",
row.len()
)));
}
}
let n_u32 = u32::try_from(n).map_err(|_| {
IgraphError::InvalidArgument("nearest_neighbor_graph: too many points".into())
})?;
let radius = if cutoff < 0.0 { f64::INFINITY } else { cutoff };
let threshold = match metric {
DistanceMetric::Euclidean => radius * radius,
DistanceMetric::Manhattan => radius,
};
let neighbor_count = if k < 0 {
n
} else {
usize::try_from(k).unwrap_or(usize::MAX)
};
let mut directed_edges: Vec<(u32, u32)> = Vec::new();
let mut cands: Vec<(f64, usize)> = Vec::new();
for a in 0..n {
cands.clear();
for b in 0..n {
if a == b {
continue;
}
let dist = match metric {
DistanceMetric::Euclidean => points[a]
.iter()
.zip(points[b].iter())
.map(|(&x, &y)| (x - y) * (x - y))
.sum::<f64>(),
DistanceMetric::Manhattan => points[a]
.iter()
.zip(points[b].iter())
.map(|(&x, &y)| (x - y).abs())
.sum::<f64>(),
};
if dist < threshold {
cands.push((dist, b));
}
}
cands.sort_by(|p, q| p.0.total_cmp(&q.0).then(p.1.cmp(&q.1)));
let take = neighbor_count.min(cands.len());
let u = u32::try_from(a)
.map_err(|_| IgraphError::Internal("nearest_neighbor_graph: vertex id overflow"))?;
for &(_, b) in &cands[..take] {
let v = u32::try_from(b)
.map_err(|_| IgraphError::Internal("nearest_neighbor_graph: vertex id overflow"))?;
directed_edges.push((u, v));
}
}
if directed {
let mut graph = Graph::new(n_u32, true)?;
graph.add_edges(directed_edges)?;
Ok(graph)
} else {
let mut seen: std::collections::HashSet<(u32, u32)> = std::collections::HashSet::new();
let mut undirected_edges: Vec<(u32, u32)> = Vec::new();
for (u, v) in directed_edges {
let key = (u.min(v), u.max(v));
if seen.insert(key) {
undirected_edges.push(key);
}
}
let mut graph = Graph::new(n_u32, false)?;
graph.add_edges(undirected_edges)?;
Ok(graph)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn arc_set(g: &Graph) -> Vec<(u32, u32)> {
let mut edges: Vec<(u32, u32)> = (0..g.ecount())
.map(|e| {
g.edge(u32::try_from(e).expect("edge id fits in u32"))
.expect("edge")
})
.collect();
edges.sort_unstable();
edges
}
fn undirected_edge_set(g: &Graph) -> Vec<(u32, u32)> {
let mut edges: Vec<(u32, u32)> = (0..g.ecount())
.map(|e| {
let (u, v) = g
.edge(u32::try_from(e).expect("edge id fits in u32"))
.expect("edge");
(u.min(v), u.max(v))
})
.collect();
edges.sort_unstable();
edges
}
#[test]
fn empty_point_set() {
let g =
nearest_neighbor_graph(&[], DistanceMetric::Euclidean, 2, 5.0, true).expect("empty");
assert_eq!(g.vcount(), 0);
assert_eq!(g.ecount(), 0);
assert!(g.is_directed());
}
#[test]
fn single_point_no_edges() {
let g = nearest_neighbor_graph(
&[vec![1.0, 6.0, 4.0]],
DistanceMetric::Euclidean,
-1,
100.0,
true,
)
.expect("singleton");
assert_eq!(g.vcount(), 1);
assert_eq!(g.ecount(), 0);
}
#[test]
fn zero_dimensional_error() {
let err = nearest_neighbor_graph(&[vec![]], DistanceMetric::Euclidean, 1, 10.0, true)
.unwrap_err();
assert!(matches!(err, IgraphError::InvalidArgument(_)));
}
#[test]
fn inconsistent_dimensions_error() {
let err = nearest_neighbor_graph(
&[vec![0.0, 0.0], vec![1.0]],
DistanceMetric::Euclidean,
1,
10.0,
true,
)
.unwrap_err();
assert!(matches!(err, IgraphError::InvalidArgument(_)));
}
#[test]
fn zero_neighbors_is_edgeless() {
let pts = vec![vec![12.0, 8.0], vec![8.0, 6.0], vec![5.0, 12.0]];
let g =
nearest_neighbor_graph(&pts, DistanceMetric::Euclidean, 0, 5.0, true).expect("zero k");
assert_eq!(g.ecount(), 0);
}
#[test]
fn zero_cutoff_is_edgeless() {
let pts = vec![vec![12.0, 8.0], vec![8.0, 6.0], vec![5.0, 12.0]];
let g = nearest_neighbor_graph(&pts, DistanceMetric::Euclidean, -1, 0.0, true)
.expect("zero cut");
assert_eq!(g.ecount(), 0);
}
#[test]
fn strict_cutoff_excludes_boundary() {
let pts = vec![vec![8.0], vec![5.0]];
let g =
nearest_neighbor_graph(&pts, DistanceMetric::Euclidean, -1, 3.0, true).expect("strict");
assert_eq!(g.ecount(), 0);
let g2 =
nearest_neighbor_graph(&pts, DistanceMetric::Euclidean, -1, 3.5, true).expect("loose");
assert_eq!(arc_set(&g2), vec![(0, 1), (1, 0)]);
}
#[test]
fn c_anchor_2d_k2_cutoff5_l2() {
let pts = vec![
vec![12.0, 8.0],
vec![8.0, 6.0],
vec![5.0, 12.0],
vec![10.0, 1.0],
vec![12.0, 2.0],
];
let g = nearest_neighbor_graph(&pts, DistanceMetric::Euclidean, 2, 5.0, true).expect("2d");
assert_eq!(arc_set(&g), vec![(0, 1), (1, 0), (3, 4), (4, 3)]);
}
#[test]
fn c_anchor_3d_k2_cutoff_inf_l2() {
let pts = vec![
vec![1.0, 6.0, 4.0],
vec![6.0, 2.0, 3.0],
vec![3.0, 6.0, 6.0],
vec![3.0, 2.0, 2.0],
vec![2.0, 3.0, 3.0],
];
let g = nearest_neighbor_graph(&pts, DistanceMetric::Euclidean, 2, -1.0, true).expect("3d");
assert_eq!(
arc_set(&g),
vec![
(0, 2),
(0, 4),
(1, 3),
(1, 4),
(2, 0),
(2, 4),
(3, 1),
(3, 4),
(4, 0),
(4, 3),
]
);
}
#[test]
fn c_anchor_4d_k1_cutoff_inf_l2() {
let pts = vec![
vec![1.0, 6.0, 4.0, 4.0],
vec![6.0, 2.0, 3.0, 3.0],
vec![3.0, 6.0, 6.0, 5.0],
vec![3.0, 2.0, 2.0, 3.0],
vec![2.0, 3.0, 3.0, 4.0],
];
let g = nearest_neighbor_graph(&pts, DistanceMetric::Euclidean, 1, -1.0, true).expect("4d");
assert_eq!(arc_set(&g), vec![(0, 2), (1, 3), (2, 0), (3, 4), (4, 3)]);
}
#[test]
fn c_anchor_2d_k2_cutoff5_l1() {
let pts = vec![
vec![12.0, 8.0],
vec![8.0, 6.0],
vec![5.0, 12.0],
vec![10.0, 1.0],
vec![12.0, 2.0],
];
let g = nearest_neighbor_graph(&pts, DistanceMetric::Manhattan, 2, 5.0, true).expect("l1");
assert_eq!(arc_set(&g), vec![(3, 4), (4, 3)]);
}
#[test]
fn undirected_collapses_reciprocal_arcs() {
let pts = vec![
vec![12.0, 8.0],
vec![8.0, 6.0],
vec![5.0, 12.0],
vec![10.0, 1.0],
vec![12.0, 2.0],
];
let g =
nearest_neighbor_graph(&pts, DistanceMetric::Euclidean, 2, 5.0, false).expect("undir");
assert!(!g.is_directed());
assert_eq!(undirected_edge_set(&g), vec![(0, 1), (3, 4)]);
}
#[test]
fn unlimited_neighbors_infinite_cutoff_is_complete() {
let pts = vec![vec![0.0], vec![1.0], vec![2.0]];
let g = nearest_neighbor_graph(&pts, DistanceMetric::Euclidean, -1, -1.0, true)
.expect("complete");
assert_eq!(g.ecount(), 6);
}
}