fn centroid_decomposition_dfs(
parents: &[usize],
vs: &[usize],
f: &mut impl FnMut(&[usize], &[usize], usize, usize),
)Examples found in repository?
crates/competitive/src/tree/centroid_decomposition.rs (line 93)
4fn centroid_decomposition_dfs(
5 parents: &[usize],
6 vs: &[usize],
7 f: &mut impl FnMut(&[usize], &[usize], usize, usize),
8) {
9 let n = vs.len();
10 assert!(n >= 2);
11 if n == 2 {
12 return;
13 }
14 let mut size = vec![1; n];
15 let mut c = !0;
16 for i in (0..n).rev() {
17 if size[i] >= n.div_ceil(2) {
18 c = i;
19 break;
20 }
21 size[parents[i]] += size[i];
22 }
23 let mut color = vec![!0u8; n];
24 let mut order = vec![!0usize; n];
25 order[c] = 0;
26 let mut count = 1usize;
27 let mut sum_size = 0usize;
28 for u in 1..n {
29 if parents[u] == c && sum_size + size[u] <= (n - 1) / 2 {
30 sum_size += size[u];
31 color[u] = 0;
32 order[u] = count;
33 count += 1;
34 }
35 }
36 for u in 1..n {
37 if color[parents[u]] == 0 {
38 color[u] = 0;
39 order[u] = count;
40 count += 1;
41 }
42 }
43 let lsize = count - 1;
44 {
45 let mut u = parents[c];
46 while u != !0 {
47 color[u] = 1;
48 order[u] = count;
49 count += 1;
50 u = parents[u];
51 }
52 }
53 for u in 0..n {
54 if u != c && color[u] == !0 {
55 color[u] = 1;
56 order[u] = count;
57 count += 1;
58 }
59 }
60 assert_eq!(count, n);
61 let rsize = n - lsize - 1;
62 let mut cparents = vec![!0usize; n];
63 let mut cvs = vec![!0usize; n];
64 let mut lparents = vec![!0usize; lsize + 1];
65 let mut lvs = vec![!0usize; lsize + 1];
66 let mut rparents = vec![!0usize; rsize + 1];
67 let mut rvs = vec![!0usize; rsize + 1];
68 for u in 0..n {
69 let i = order[u];
70 cvs[i] = vs[u];
71 if color[u] != 1 {
72 lvs[i] = vs[u];
73 }
74 if color[u] != 0 {
75 rvs[if i == 0 { 0 } else { i - lsize }] = vs[u];
76 }
77 }
78 for u in 1..n {
79 let mut x = order[u];
80 let mut y = order[parents[u]];
81 if x > y {
82 swap(&mut x, &mut y);
83 }
84 cparents[y] = x;
85 if color[u] != 1 && color[parents[u]] != 1 {
86 lparents[y] = x;
87 }
88 if color[u] != 0 && color[parents[u]] != 0 {
89 rparents[if y == 0 { 0 } else { y - lsize }] = if x == 0 { 0 } else { x - lsize };
90 }
91 }
92 f(&cparents, &cvs, lsize, rsize);
93 centroid_decomposition_dfs(&lparents, &lvs, f);
94 centroid_decomposition_dfs(&rparents, &rvs, f);
95}
96
97impl UndirectedSparseGraph {
98 /// 1/3 centroid decomposition
99 ///
100 /// - f: (parents: &[usize], vs: &[usize], lsize: usize, rsize: usize)
101 /// - 0: root, 1..=lsize: left subtree, lsize+1..=lsize+rsize: right subtree
102 pub fn centroid_decomposition(&self, mut f: impl FnMut(&[usize], &[usize], usize, usize)) {
103 let n = self.vertices_size();
104 if n <= 1 {
105 return;
106 }
107 let mut vs = Vec::with_capacity(n);
108 let mut parents = vec![!0; n];
109 vs.push(0usize);
110 for l in 0..n {
111 let u = vs[l];
112 for a in self.adjacencies(u) {
113 if a.to != parents[u] {
114 vs.push(a.to);
115 parents[a.to] = u;
116 }
117 }
118 }
119 let mut new_idx = vec![0; n];
120 for i in 0..n {
121 new_idx[vs[i]] = i;
122 }
123 let mut new_parent = vec![!0; n];
124 for i in 1..n {
125 new_parent[new_idx[i]] = new_idx[parents[i]];
126 }
127 centroid_decomposition_dfs(&new_parent, &vs, &mut f);
128 }