use super::csr::CsrGraph;
use crate::rng::Xorshift64;
pub(super) struct CoarseningLevel {
pub(super) graph: CsrGraph,
pub(super) mapping: Vec<u32>,
}
pub(super) fn coarsen_one_level(
graph: &CsrGraph,
min_vertices: usize,
rng: &mut Xorshift64,
part: Option<&[u8]>,
) -> Option<CoarseningLevel> {
crate::meter::charge(graph.pass_units());
let n = graph.num_vertices();
if n <= min_vertices {
return None;
}
let mut perm: Vec<usize> = (0..n).collect();
perm.sort_by_key(|&v| {
let degree = graph.offsets[v + 1] - graph.offsets[v];
(degree, graph.vertex_weights[v])
});
let mut i = 0;
while i < n {
let mut j = i + 1;
let degree = graph.offsets[perm[i] + 1] - graph.offsets[perm[i]];
while j < n && (graph.offsets[perm[j] + 1] - graph.offsets[perm[j]]) == degree {
j += 1;
}
for k in (i + 1..j).rev() {
let l = i + (rng.next_u64() as usize) % (k - i + 1);
perm.swap(k, l);
}
i = j;
}
let mut match_of = vec![None; n];
let mut coarse_id: Vec<u32> = vec![0; n];
let mut num_coarse: u32 = 0;
for &vtx in &perm {
if match_of[vtx].is_some() {
continue;
}
let start = graph.offsets[vtx] as usize;
let end = graph.offsets[vtx + 1] as usize;
let neighbors = &graph.neighbors[start..end];
let weights = &graph.edge_weights[start..end];
let mut best_neighbor = None;
let mut best_wt: u32 = 0;
let mut best_same_part = false;
for (&nb, &w) in neighbors.iter().zip(weights) {
if match_of[nb as usize].is_some() {
continue;
}
let same_part = part.is_some_and(|p| p[nb as usize] == p[vtx]);
if same_part && !best_same_part {
best_wt = w;
best_neighbor = Some(nb);
best_same_part = true;
} else if same_part == best_same_part && w > best_wt {
best_wt = w;
best_neighbor = Some(nb);
}
}
if let Some(neighbor) = best_neighbor {
match_of[vtx] = Some(neighbor);
match_of[neighbor as usize] = Some(vtx as u32);
coarse_id[vtx] = num_coarse;
coarse_id[neighbor as usize] = num_coarse;
num_coarse += 1;
} else {
match_of[vtx] = Some(vtx as u32);
coarse_id[vtx] = num_coarse;
num_coarse += 1;
}
}
let nc = num_coarse as usize;
if nc >= n * 9 / 10 {
return None; }
let mut coarse_vwgt = vec![0u32; nc];
for v in 0..n {
coarse_vwgt[coarse_id[v] as usize] += graph.vertex_weights[v];
}
let mut counts: Vec<u32> = vec![0; nc];
for v in 0..n {
let cv = coarse_id[v] as usize;
for &nb in graph.neighbors(v) {
let cnb = coarse_id[nb as usize];
if cnb as usize != cv {
counts[cv] += 1;
}
}
}
let mut offsets: Vec<u32> = Vec::with_capacity(nc + 1);
offsets.push(0);
let mut acc = 0u32;
for &c in &counts {
acc += c;
offsets.push(acc);
}
let total = acc as usize;
let mut flat: Vec<(u32, u32)> = vec![(0, 0); total];
let mut cur = counts; for c in cur.iter_mut() {
*c = 0;
}
for v in 0..n {
let cv = coarse_id[v] as usize;
let start = graph.offsets[v] as usize;
let end = graph.offsets[v + 1] as usize;
let neighbors = &graph.neighbors[start..end];
let weights = &graph.edge_weights[start..end];
for (&nb, &w) in neighbors.iter().zip(weights) {
let cnb = coarse_id[nb as usize];
if cnb as usize != cv {
let pos = (offsets[cv] + cur[cv]) as usize;
flat[pos] = (cnb, w);
cur[cv] += 1;
}
}
}
let mut coarse_offsets = Vec::with_capacity(nc + 1);
let mut coarse_neighbors = Vec::new();
let mut coarse_edge_weights = Vec::new();
coarse_offsets.push(0u32);
for v in 0..nc {
let start = offsets[v] as usize;
let end = offsets[v + 1] as usize;
let slice = &mut flat[start..end];
slice.sort_unstable_by_key(|&(nb, _)| nb);
let mut i = 0;
while i < slice.len() {
let nb = slice[i].0;
let mut wt = 0u32;
while i < slice.len() && slice[i].0 == nb {
wt = wt
.checked_add(slice[i].1)
.expect("a coarse edge cannot outweigh all fine edges");
i += 1;
}
coarse_neighbors.push(nb);
coarse_edge_weights.push(wt);
}
coarse_offsets.push(coarse_neighbors.len() as u32);
}
Some(CoarseningLevel {
graph: CsrGraph {
offsets: coarse_offsets,
neighbors: coarse_neighbors,
vertex_weights: coarse_vwgt,
edge_weights: coarse_edge_weights,
},
mapping: coarse_id,
})
}