Skip to main content

oxirs_graphrag/
graph_partitioner.rs

1//! Graph partitioning using greedy and label-propagation methods.
2//!
3//! Provides k-way graph partitioning for knowledge graphs, useful for
4//! distributed SPARQL query planning, parallel community processing,
5//! and load-balanced federated queries.
6
7use std::collections::HashMap;
8
9// ─────────────────────────────────────────────────────────────────────────────
10// Types
11// ─────────────────────────────────────────────────────────────────────────────
12
13/// Assignment of a single node to a partition.
14#[derive(Debug, Clone, PartialEq, Eq)]
15pub struct GraphPartition {
16    /// Node identifier (matches the input `nodes` slice).
17    pub node_id: String,
18    /// Zero-based partition index (0..k).
19    pub partition: usize,
20}
21
22/// Result of a graph partitioning run.
23#[derive(Debug, Clone)]
24pub struct PartitionResult {
25    /// Per-node partition assignments.
26    pub assignments: Vec<GraphPartition>,
27    /// Number of partitions k.
28    pub num_partitions: usize,
29    /// Number of edges whose endpoints are in different partitions.
30    pub cut_edges: usize,
31    /// Balance score in [0, 1]. 1.0 = perfectly balanced.
32    pub balance_score: f64,
33}
34
35/// Available partitioning algorithms.
36#[derive(Debug, Clone, PartialEq, Eq)]
37pub enum PartitionMethod {
38    /// Round-robin greedy assignment followed by a local-improvement step.
39    Greedy,
40    /// Iterative label propagation (majority vote of neighbour labels).
41    LabelPropagation,
42    /// Recursive bisection by BFS-cut (produces a binary-tree of cuts).
43    Bisection,
44}
45
46// ─────────────────────────────────────────────────────────────────────────────
47// GraphPartitioner
48// ─────────────────────────────────────────────────────────────────────────────
49
50/// Configures and runs graph partitioning.
51#[derive(Debug, Clone)]
52pub struct GraphPartitioner {
53    /// Number of partitions k.
54    pub num_partitions: usize,
55    /// Algorithm selection.
56    pub method: PartitionMethod,
57    /// Maximum iterations (used by LabelPropagation and improvement loops).
58    pub max_iterations: usize,
59}
60
61impl GraphPartitioner {
62    /// Create a partitioner with the default method (Greedy) and 20 iterations.
63    pub fn new(num_partitions: usize) -> Self {
64        Self {
65            num_partitions: num_partitions.max(1),
66            method: PartitionMethod::Greedy,
67            max_iterations: 20,
68        }
69    }
70
71    /// Select the partitioning method.
72    pub fn with_method(mut self, method: PartitionMethod) -> Self {
73        self.method = method;
74        self
75    }
76
77    /// Override the maximum iteration count.
78    pub fn with_max_iterations(mut self, max_iter: usize) -> Self {
79        self.max_iterations = max_iter;
80        self
81    }
82
83    /// Partition the graph and return a `PartitionResult`.
84    ///
85    /// # Arguments
86    /// * `nodes` – list of node identifiers (must be unique).
87    /// * `edges` – list of `(from, to)` string pairs; direction is ignored.
88    pub fn partition(&self, nodes: &[String], edges: &[(String, String)]) -> PartitionResult {
89        if nodes.is_empty() {
90            return PartitionResult {
91                assignments: vec![],
92                num_partitions: self.num_partitions,
93                cut_edges: 0,
94                balance_score: 1.0,
95            };
96        }
97
98        let k = self.num_partitions;
99
100        let labels = match &self.method {
101            PartitionMethod::Greedy => Self::greedy_partition(nodes, edges, k),
102            PartitionMethod::LabelPropagation => {
103                Self::label_propagation(nodes, edges, k, self.max_iterations)
104            }
105            PartitionMethod::Bisection => Self::bisection_partition(nodes, edges, k),
106        };
107
108        // Map node-string → index for cut-edge counting
109        let node_idx: HashMap<&str, usize> = nodes
110            .iter()
111            .enumerate()
112            .map(|(i, n)| (n.as_str(), i))
113            .collect();
114
115        let int_edges: Vec<(usize, usize)> = edges
116            .iter()
117            .filter_map(|(a, b)| {
118                let ai = node_idx.get(a.as_str())?;
119                let bi = node_idx.get(b.as_str())?;
120                Some((*ai, *bi))
121            })
122            .collect();
123
124        let cut_edges = Self::count_cut_edges(&labels, &int_edges);
125        let balance = Self::balance_score(&labels, k);
126
127        let assignments = nodes
128            .iter()
129            .enumerate()
130            .map(|(i, n)| GraphPartition {
131                node_id: n.clone(),
132                partition: labels[i],
133            })
134            .collect();
135
136        PartitionResult {
137            assignments,
138            num_partitions: k,
139            cut_edges,
140            balance_score: balance,
141        }
142    }
143
144    // ── Algorithm implementations ─────────────────────────────────────────────
145
146    /// Greedy round-robin assignment with a single local-improvement sweep.
147    ///
148    /// Returns a per-node partition vector (indices into `nodes`).
149    pub fn greedy_partition(nodes: &[String], edges: &[(String, String)], k: usize) -> Vec<usize> {
150        let k = k.max(1);
151        let n = nodes.len();
152        if n == 0 {
153            return vec![];
154        }
155
156        // Initial round-robin assignment
157        let mut labels: Vec<usize> = (0..n).map(|i| i % k).collect();
158
159        // Build adjacency list
160        let adj = Self::build_adjacency(nodes, edges);
161
162        // One balance-constrained local-improvement pass: for each node, consider
163        // moving it to the partition most common among its neighbours, but only
164        // if the destination partition is not larger than a balanced target.
165        let target_per_partition = (n + k - 1) / k; // ceiling division
166        for i in 0..n {
167            if adj[i].is_empty() {
168                continue;
169            }
170            let mut counts = vec![0usize; k];
171            for &nb in &adj[i] {
172                counts[labels[nb]] += 1;
173            }
174            // Find best neighbour-majority partition that is not over-full
175            let current_part = labels[i];
176            let mut best_part = current_part;
177            let mut best_count = counts[current_part];
178
179            // Count current partition sizes
180            let mut part_sizes = vec![0usize; k];
181            for &l in labels.iter() {
182                part_sizes[l.min(k - 1)] += 1;
183            }
184
185            for (p, &c) in counts.iter().enumerate() {
186                if c > best_count && part_sizes[p] < target_per_partition {
187                    best_part = p;
188                    best_count = c;
189                }
190            }
191            labels[i] = best_part;
192        }
193        labels
194    }
195
196    /// Iterative label propagation.
197    ///
198    /// Each node adopts the most frequent label among its neighbours.
199    /// Labels are initialised in round-robin fashion and the propagation
200    /// continues for at most `max_iter` rounds or until convergence.
201    pub fn label_propagation(
202        nodes: &[String],
203        edges: &[(String, String)],
204        k: usize,
205        max_iter: usize,
206    ) -> Vec<usize> {
207        let k = k.max(1);
208        let n = nodes.len();
209        if n == 0 {
210            return vec![];
211        }
212
213        let adj = Self::build_adjacency(nodes, edges);
214        // Round-robin init
215        let mut labels: Vec<usize> = (0..n).map(|i| i % k).collect();
216
217        for _ in 0..max_iter {
218            let mut changed = false;
219            let prev = labels.clone();
220
221            for i in 0..n {
222                if adj[i].is_empty() {
223                    continue;
224                }
225                let mut counts = vec![0usize; k];
226                for &nb in &adj[i] {
227                    counts[prev[nb]] += 1;
228                }
229                // Find label with max count; break ties by smallest label index
230                let best = counts
231                    .iter()
232                    .enumerate()
233                    .max_by(|(la, &ca), (lb, &cb)| ca.cmp(&cb).then(lb.cmp(la)))
234                    .map(|(p, _)| p)
235                    .unwrap_or(labels[i]);
236
237                if best != labels[i] {
238                    labels[i] = best;
239                    changed = true;
240                }
241            }
242
243            if !changed {
244                break;
245            }
246        }
247
248        // Ensure labels stay in [0, k)
249        for l in &mut labels {
250            *l = (*l).min(k - 1);
251        }
252        labels
253    }
254
255    /// Recursive bisection: repeatedly split the largest partition by BFS.
256    pub fn bisection_partition(
257        nodes: &[String],
258        edges: &[(String, String)],
259        k: usize,
260    ) -> Vec<usize> {
261        let k = k.max(1);
262        let n = nodes.len();
263        if n == 0 {
264            return vec![];
265        }
266
267        let adj = Self::build_adjacency(nodes, edges);
268        let mut labels = vec![0usize; n];
269
270        // We produce up to k partitions by splitting iteratively
271        let target_splits = k.saturating_sub(1);
272        let mut current_k = 1usize;
273
274        for _ in 0..target_splits {
275            if current_k >= k {
276                break;
277            }
278            // Find the partition with the most nodes
279            let mut part_sizes = vec![0usize; current_k];
280            for &l in &labels {
281                part_sizes[l] += 1;
282            }
283            let largest_part = part_sizes
284                .iter()
285                .enumerate()
286                .max_by_key(|(_, &s)| s)
287                .map(|(p, _)| p)
288                .unwrap_or(0);
289
290            // Collect nodes in that partition
291            let part_nodes: Vec<usize> = (0..n).filter(|&i| labels[i] == largest_part).collect();
292            if part_nodes.len() < 2 {
293                break;
294            }
295
296            // BFS from the first node; assign second half to new_part
297            let half = part_nodes.len() / 2;
298            let new_part = current_k;
299            current_k += 1;
300
301            let mut bfs_order: Vec<usize> = Vec::with_capacity(part_nodes.len());
302            let mut visited = vec![false; n];
303            let mut queue = std::collections::VecDeque::new();
304            let start = part_nodes[0];
305            queue.push_back(start);
306            visited[start] = true;
307
308            while let Some(node) = queue.pop_front() {
309                if labels[node] == largest_part {
310                    bfs_order.push(node);
311                }
312                for &nb in &adj[node] {
313                    if !visited[nb] && labels[nb] == largest_part {
314                        visited[nb] = true;
315                        queue.push_back(nb);
316                    }
317                }
318            }
319
320            // Any nodes not reached by BFS are also in part_nodes
321            for &pn in &part_nodes {
322                if !bfs_order.contains(&pn) {
323                    bfs_order.push(pn);
324                }
325            }
326
327            // Move second half to new partition
328            for &node in bfs_order.iter().skip(half) {
329                labels[node] = new_part;
330            }
331        }
332
333        labels
334    }
335
336    // ── Utility functions ─────────────────────────────────────────────────────
337
338    /// Count edges whose endpoints are in different partitions.
339    pub fn count_cut_edges(assignments: &[usize], edges: &[(usize, usize)]) -> usize {
340        edges
341            .iter()
342            .filter(|&&(a, b)| {
343                a < assignments.len() && b < assignments.len() && assignments[a] != assignments[b]
344            })
345            .count()
346    }
347
348    /// Balance score in [0, 1].  1.0 = perfectly balanced.
349    ///
350    /// Defined as `min_size / max_size` where sizes are partition cardinalities.
351    /// Returns 1.0 for empty assignments or k=1.
352    pub fn balance_score(assignments: &[usize], k: usize) -> f64 {
353        if assignments.is_empty() || k <= 1 {
354            return 1.0;
355        }
356        let mut counts = vec![0usize; k];
357        for &l in assignments {
358            let idx = l.min(k - 1);
359            counts[idx] += 1;
360        }
361        let max = *counts.iter().max().unwrap_or(&0);
362        let min = *counts.iter().min().unwrap_or(&0);
363        if max == 0 {
364            return 1.0;
365        }
366        min as f64 / max as f64
367    }
368
369    /// Build an undirected adjacency list indexed by position in `nodes`.
370    pub fn build_adjacency(nodes: &[String], edges: &[(String, String)]) -> Vec<Vec<usize>> {
371        let n = nodes.len();
372        let node_idx: HashMap<&str, usize> = nodes
373            .iter()
374            .enumerate()
375            .map(|(i, s)| (s.as_str(), i))
376            .collect();
377
378        let mut adj = vec![vec![]; n];
379        for (a, b) in edges {
380            if let (Some(&ai), Some(&bi)) = (node_idx.get(a.as_str()), node_idx.get(b.as_str())) {
381                if ai != bi {
382                    adj[ai].push(bi);
383                    adj[bi].push(ai);
384                }
385            }
386        }
387        adj
388    }
389}
390
391// ─────────────────────────────────────────────────────────────────────────────
392// Tests
393// ─────────────────────────────────────────────────────────────────────────────
394
395#[cfg(test)]
396mod tests {
397    use super::*;
398
399    fn make_nodes(n: usize) -> Vec<String> {
400        (0..n).map(|i| format!("node_{i}")).collect()
401    }
402
403    fn chain_edges(n: usize) -> Vec<(String, String)> {
404        (0..n.saturating_sub(1))
405            .map(|i| (format!("node_{i}"), format!("node_{}", i + 1)))
406            .collect()
407    }
408
409    // ── GraphPartitioner::new ─────────────────────────────────────────────────
410
411    #[test]
412    fn test_new_default_method() {
413        let gp = GraphPartitioner::new(4);
414        assert_eq!(gp.num_partitions, 4);
415        assert_eq!(gp.method, PartitionMethod::Greedy);
416        assert_eq!(gp.max_iterations, 20);
417    }
418
419    #[test]
420    fn test_new_zero_becomes_one() {
421        let gp = GraphPartitioner::new(0);
422        assert_eq!(gp.num_partitions, 1);
423    }
424
425    #[test]
426    fn test_with_method_label_propagation() {
427        let gp = GraphPartitioner::new(3).with_method(PartitionMethod::LabelPropagation);
428        assert_eq!(gp.method, PartitionMethod::LabelPropagation);
429    }
430
431    #[test]
432    fn test_with_max_iterations() {
433        let gp = GraphPartitioner::new(3).with_max_iterations(50);
434        assert_eq!(gp.max_iterations, 50);
435    }
436
437    // ── partition – empty input ────────────────────────────────────────────────
438
439    #[test]
440    fn test_partition_empty_nodes() {
441        let gp = GraphPartitioner::new(3);
442        let result = gp.partition(&[], &[]);
443        assert!(result.assignments.is_empty());
444        assert_eq!(result.cut_edges, 0);
445        assert_eq!(result.balance_score, 1.0);
446    }
447
448    // ── partition – single node ───────────────────────────────────────────────
449
450    #[test]
451    fn test_partition_single_node() {
452        let gp = GraphPartitioner::new(3);
453        let nodes = vec!["A".to_string()];
454        let result = gp.partition(&nodes, &[]);
455        assert_eq!(result.assignments.len(), 1);
456        assert_eq!(result.assignments[0].node_id, "A");
457        assert_eq!(result.cut_edges, 0);
458    }
459
460    // ── partition – correct assignment count ──────────────────────────────────
461
462    #[test]
463    fn test_partition_returns_all_nodes() {
464        let nodes = make_nodes(10);
465        let edges = chain_edges(10);
466        let gp = GraphPartitioner::new(3);
467        let result = gp.partition(&nodes, &edges);
468        assert_eq!(result.assignments.len(), 10);
469    }
470
471    #[test]
472    fn test_partition_labels_in_range() {
473        let nodes = make_nodes(12);
474        let edges = chain_edges(12);
475        let k = 4;
476        let gp = GraphPartitioner::new(k);
477        let result = gp.partition(&nodes, &edges);
478        for a in &result.assignments {
479            assert!(a.partition < k, "label {} out of range", a.partition);
480        }
481    }
482
483    #[test]
484    fn test_partition_num_partitions_field() {
485        let nodes = make_nodes(6);
486        let gp = GraphPartitioner::new(3);
487        let result = gp.partition(&nodes, &[]);
488        assert_eq!(result.num_partitions, 3);
489    }
490
491    // ── greedy_partition ──────────────────────────────────────────────────────
492
493    #[test]
494    fn test_greedy_partition_count() {
495        let nodes = make_nodes(9);
496        let edges = chain_edges(9);
497        let labels = GraphPartitioner::greedy_partition(&nodes, &edges, 3);
498        assert_eq!(labels.len(), 9);
499    }
500
501    #[test]
502    fn test_greedy_partition_labels_valid() {
503        let nodes = make_nodes(9);
504        let edges = chain_edges(9);
505        let labels = GraphPartitioner::greedy_partition(&nodes, &edges, 3);
506        for &l in &labels {
507            assert!(l < 3);
508        }
509    }
510
511    #[test]
512    fn test_greedy_partition_empty() {
513        let labels = GraphPartitioner::greedy_partition(&[], &[], 3);
514        assert!(labels.is_empty());
515    }
516
517    #[test]
518    fn test_greedy_partition_k1() {
519        let nodes = make_nodes(5);
520        let labels = GraphPartitioner::greedy_partition(&nodes, &[], 1);
521        assert!(labels.iter().all(|&l| l == 0));
522    }
523
524    // ── label_propagation ─────────────────────────────────────────────────────
525
526    #[test]
527    fn test_label_propagation_count() {
528        let nodes = make_nodes(8);
529        let edges = chain_edges(8);
530        let labels = GraphPartitioner::label_propagation(&nodes, &edges, 2, 10);
531        assert_eq!(labels.len(), 8);
532    }
533
534    #[test]
535    fn test_label_propagation_labels_valid() {
536        let nodes = make_nodes(8);
537        let edges = chain_edges(8);
538        let labels = GraphPartitioner::label_propagation(&nodes, &edges, 4, 10);
539        for &l in &labels {
540            assert!(l < 4);
541        }
542    }
543
544    #[test]
545    fn test_label_propagation_empty() {
546        let labels = GraphPartitioner::label_propagation(&[], &[], 3, 10);
547        assert!(labels.is_empty());
548    }
549
550    #[test]
551    fn test_label_propagation_converges() {
552        // With many iterations the result should still be valid
553        let nodes = make_nodes(10);
554        let edges = chain_edges(10);
555        let labels = GraphPartitioner::label_propagation(&nodes, &edges, 3, 100);
556        assert_eq!(labels.len(), 10);
557        for &l in &labels {
558            assert!(l < 3);
559        }
560    }
561
562    // ── count_cut_edges ───────────────────────────────────────────────────────
563
564    #[test]
565    fn test_count_cut_edges_none() {
566        // All in partition 0
567        let assignments = vec![0, 0, 0];
568        let edges = vec![(0, 1), (1, 2)];
569        assert_eq!(GraphPartitioner::count_cut_edges(&assignments, &edges), 0);
570    }
571
572    #[test]
573    fn test_count_cut_edges_all() {
574        // Alternating partitions
575        let assignments = vec![0, 1, 0, 1];
576        let edges = vec![(0, 1), (1, 2), (2, 3)];
577        assert_eq!(GraphPartitioner::count_cut_edges(&assignments, &edges), 3);
578    }
579
580    #[test]
581    fn test_count_cut_edges_empty() {
582        assert_eq!(GraphPartitioner::count_cut_edges(&[], &[]), 0);
583    }
584
585    #[test]
586    fn test_count_cut_edges_partial() {
587        let assignments = vec![0, 0, 1, 1];
588        let edges = vec![(0, 1), (1, 2), (2, 3)];
589        // edge (1,2) is cut
590        assert_eq!(GraphPartitioner::count_cut_edges(&assignments, &edges), 1);
591    }
592
593    // ── balance_score ─────────────────────────────────────────────────────────
594
595    #[test]
596    fn test_balance_score_perfect() {
597        let assignments = vec![0, 1, 0, 1];
598        let score = GraphPartitioner::balance_score(&assignments, 2);
599        assert!((score - 1.0).abs() < 1e-10);
600    }
601
602    #[test]
603    fn test_balance_score_empty() {
604        assert!((GraphPartitioner::balance_score(&[], 3) - 1.0).abs() < 1e-10);
605    }
606
607    #[test]
608    fn test_balance_score_k1() {
609        let assignments = vec![0, 0, 0];
610        assert!((GraphPartitioner::balance_score(&assignments, 1) - 1.0).abs() < 1e-10);
611    }
612
613    #[test]
614    fn test_balance_score_imbalanced() {
615        // 1 node in part 0, 3 in part 1 → min/max = 1/3
616        let assignments = vec![0, 1, 1, 1];
617        let score = GraphPartitioner::balance_score(&assignments, 2);
618        assert!((score - 1.0 / 3.0).abs() < 1e-10);
619    }
620
621    #[test]
622    fn test_balance_score_in_range() {
623        let assignments = vec![0, 1, 2, 0, 1, 2, 0];
624        let score = GraphPartitioner::balance_score(&assignments, 3);
625        assert!((0.0..=1.0).contains(&score));
626    }
627
628    // ── build_adjacency ───────────────────────────────────────────────────────
629
630    #[test]
631    fn test_build_adjacency_empty() {
632        let adj = GraphPartitioner::build_adjacency(&[], &[]);
633        assert!(adj.is_empty());
634    }
635
636    #[test]
637    fn test_build_adjacency_chain() {
638        let nodes = make_nodes(3);
639        let edges = chain_edges(3);
640        let adj = GraphPartitioner::build_adjacency(&nodes, &edges);
641        assert_eq!(adj.len(), 3);
642        assert!(adj[0].contains(&1));
643        assert!(adj[1].contains(&0));
644        assert!(adj[1].contains(&2));
645        assert!(adj[2].contains(&1));
646    }
647
648    #[test]
649    fn test_build_adjacency_no_self_loops() {
650        let nodes = vec!["A".to_string()];
651        let edges = vec![("A".to_string(), "A".to_string())];
652        let adj = GraphPartitioner::build_adjacency(&nodes, &edges);
653        assert!(adj[0].is_empty());
654    }
655
656    #[test]
657    fn test_build_adjacency_unknown_node_ignored() {
658        let nodes = make_nodes(2);
659        // edge references node_99 which does not exist
660        let edges = vec![("node_0".to_string(), "node_99".to_string())];
661        let adj = GraphPartitioner::build_adjacency(&nodes, &edges);
662        assert!(adj[0].is_empty());
663    }
664
665    // ── bisection partition ───────────────────────────────────────────────────
666
667    #[test]
668    fn test_bisection_labels_valid() {
669        let nodes = make_nodes(8);
670        let edges = chain_edges(8);
671        let labels = GraphPartitioner::bisection_partition(&nodes, &edges, 4);
672        for &l in &labels {
673            assert!(l < 4);
674        }
675    }
676
677    #[test]
678    fn test_bisection_count() {
679        let nodes = make_nodes(6);
680        let labels = GraphPartitioner::bisection_partition(&nodes, &[], 3);
681        assert_eq!(labels.len(), 6);
682    }
683
684    // ── LabelPropagation method end-to-end ───────────────────────────────────
685
686    #[test]
687    fn test_partition_label_propagation_method() {
688        let nodes = make_nodes(10);
689        let edges = chain_edges(10);
690        let gp = GraphPartitioner::new(2).with_method(PartitionMethod::LabelPropagation);
691        let result = gp.partition(&nodes, &edges);
692        assert_eq!(result.assignments.len(), 10);
693        assert!(result.balance_score >= 0.0 && result.balance_score <= 1.0);
694    }
695
696    // ── Bisection method end-to-end ───────────────────────────────────────────
697
698    #[test]
699    fn test_partition_bisection_method() {
700        let nodes = make_nodes(8);
701        let edges = chain_edges(8);
702        let gp = GraphPartitioner::new(4).with_method(PartitionMethod::Bisection);
703        let result = gp.partition(&nodes, &edges);
704        assert_eq!(result.assignments.len(), 8);
705        for a in &result.assignments {
706            assert!(a.partition < 4);
707        }
708    }
709
710    // ── GraphPartition struct ─────────────────────────────────────────────────
711
712    #[test]
713    fn test_graph_partition_fields() {
714        let gp = GraphPartition {
715            node_id: "A".to_string(),
716            partition: 2,
717        };
718        assert_eq!(gp.node_id, "A");
719        assert_eq!(gp.partition, 2);
720    }
721
722    #[test]
723    fn test_graph_partition_clone() {
724        let gp = GraphPartition {
725            node_id: "X".to_string(),
726            partition: 1,
727        };
728        let gp2 = gp.clone();
729        assert_eq!(gp, gp2);
730    }
731
732    // ── Dense fully-connected graph ───────────────────────────────────────────
733
734    #[test]
735    fn test_fully_connected_partition() {
736        let nodes = make_nodes(4);
737        let edges: Vec<(String, String)> = vec![
738            ("node_0".to_string(), "node_1".to_string()),
739            ("node_0".to_string(), "node_2".to_string()),
740            ("node_0".to_string(), "node_3".to_string()),
741            ("node_1".to_string(), "node_2".to_string()),
742            ("node_1".to_string(), "node_3".to_string()),
743            ("node_2".to_string(), "node_3".to_string()),
744        ];
745        let gp = GraphPartitioner::new(2);
746        let result = gp.partition(&nodes, &edges);
747        assert_eq!(result.assignments.len(), 4);
748        // Some edges must be cut for k=2 in a fully connected graph
749        assert!(result.cut_edges > 0);
750    }
751
752    // ── PartitionResult fields ────────────────────────────────────────────────
753
754    #[test]
755    fn test_partition_result_fields() {
756        let nodes = make_nodes(6);
757        let edges = chain_edges(6);
758        let gp = GraphPartitioner::new(2);
759        let result = gp.partition(&nodes, &edges);
760        assert_eq!(result.num_partitions, 2);
761        assert!(result.balance_score >= 0.0 && result.balance_score <= 1.0);
762    }
763
764    #[test]
765    fn test_partition_result_assignment_node_ids() {
766        let nodes = make_nodes(4);
767        let gp = GraphPartitioner::new(2);
768        let result = gp.partition(&nodes, &[]);
769        let ids: Vec<&str> = result
770            .assignments
771            .iter()
772            .map(|a| a.node_id.as_str())
773            .collect();
774        assert!(ids.contains(&"node_0"));
775        assert!(ids.contains(&"node_3"));
776    }
777
778    // ── greedy with no edges ──────────────────────────────────────────────────
779
780    #[test]
781    fn test_greedy_no_edges() {
782        let nodes = make_nodes(6);
783        let labels = GraphPartitioner::greedy_partition(&nodes, &[], 3);
784        assert_eq!(labels.len(), 6);
785        for &l in &labels {
786            assert!(l < 3);
787        }
788    }
789
790    // ── label propagation with k > node count ────────────────────────────────
791
792    #[test]
793    fn test_label_propagation_k_larger_than_n() {
794        let nodes = make_nodes(3);
795        let edges = chain_edges(3);
796        let labels = GraphPartitioner::label_propagation(&nodes, &edges, 10, 5);
797        assert_eq!(labels.len(), 3);
798        for &l in &labels {
799            assert!(l < 10);
800        }
801    }
802
803    // ── balance score all same partition ─────────────────────────────────────
804
805    #[test]
806    fn test_balance_score_all_same() {
807        let assignments = vec![0, 0, 0, 0];
808        // Only partition 0 has nodes; partition 1 has 0 → min=0, max=4 → 0.0
809        let score = GraphPartitioner::balance_score(&assignments, 2);
810        assert_eq!(score, 0.0);
811    }
812
813    // ── cut edges out-of-range index ─────────────────────────────────────────
814
815    #[test]
816    fn test_count_cut_edges_out_of_range() {
817        let assignments = vec![0, 1];
818        // edge (0, 5) — index 5 is out of range, should be ignored
819        let edges = vec![(0usize, 5usize)];
820        assert_eq!(GraphPartitioner::count_cut_edges(&assignments, &edges), 0);
821    }
822
823    // ── bisection single node ─────────────────────────────────────────────────
824
825    #[test]
826    fn test_bisection_single_node() {
827        let nodes = vec!["A".to_string()];
828        let labels = GraphPartitioner::bisection_partition(&nodes, &[], 2);
829        assert_eq!(labels.len(), 1);
830        assert_eq!(labels[0], 0);
831    }
832
833    // ── PartitionMethod debug ─────────────────────────────────────────────────
834
835    #[test]
836    fn test_partition_method_debug() {
837        let m = PartitionMethod::LabelPropagation;
838        let s = format!("{m:?}");
839        assert!(s.contains("LabelPropagation"));
840    }
841}