nekolib/graph/
hld.rs

1//! HL 分解。
2
3use std::fmt::Debug;
4
5/// HL 分解。
6///
7/// <https://codeforces.com/blog/entry/53170>, Easiest HLD with subtree queries.
8///
9/// <https://atcoder.jp/contests/abc294/submissions/39923465>
10#[derive(Clone, Debug)]
11pub struct Hld {
12    perm: Vec<usize>,
13    perm_inv: Vec<usize>,
14    inout: Vec<(usize, usize)>,
15    heavy: Vec<usize>,
16    par: Vec<usize>,
17    depth: Vec<usize>,
18}
19
20#[derive(Clone, Copy, Debug, Eq, PartialEq)]
21pub enum HlEdge {
22    Heavy(usize, usize),
23    Light(usize, usize),
24}
25use HlEdge::{Heavy, Light};
26
27#[derive(Clone, Copy, Debug, Eq, PartialEq)]
28pub enum Direction {
29    Asc,
30    Desc,
31}
32use Direction::{Asc, Desc};
33
34impl HlEdge {
35    pub fn rev(self) -> Self {
36        match self {
37            Heavy(u, v) => Heavy(v, u),
38            Light(u, v) => Light(v, u),
39        }
40    }
41    pub fn inner(self) -> (usize, usize) {
42        match self {
43            Heavy(u, v) | Light(u, v) => (u, v),
44        }
45    }
46}
47
48impl Hld {
49    // g[v] は、子方向への隣接頂点のみを持つとする。
50    pub fn new(mut g: Vec<Vec<usize>>, r: usize) -> Self {
51        let n = g.len();
52        let mut size = vec![0; n];
53        Self::dfs_size(&mut g, r, &mut size);
54
55        let perm = Self::dfs_order(&g, r);
56        let perm_inv = Self::inv(&perm);
57        let g = Self::relabel(&g, &perm);
58        let (inout, heavy) = Self::dfs_hld(&g, r);
59
60        let (par, depth) = Self::par_depth(&g, r);
61        Self { perm, perm_inv, inout, heavy, par, depth }
62    }
63
64    pub fn lca_decoded(&self, ou: usize, ov: usize) -> usize {
65        let u = self.perm[ou];
66        let v = self.perm[ov];
67        let w = self.lca(u, v);
68        self.perm_inv[w]
69    }
70
71    pub fn path(
72        &self,
73        ou: usize,
74        ov: usize,
75    ) -> impl Iterator<Item = (HlEdge, Direction)> {
76        let u = self.perm[ou];
77        let v = self.perm[ov];
78        let w = self.lca(u, v);
79        let asc = self.ascend(u, w).map(|e| (e, Asc));
80        let desc = self.ascend(v, w).map(|e| (e.rev(), Desc)).rev();
81        asc.chain(desc)
82    }
83
84    pub fn subtree_range(&self, ov: usize) -> (usize, usize) {
85        self.inout[self.perm[ov]]
86    }
87
88    pub fn encode(&self, ov: usize) -> usize { self.perm[ov] }
89    pub fn decode(&self, v: usize) -> usize { self.perm_inv[v] }
90
91    fn lca(&self, u: usize, v: usize) -> usize {
92        let dh = |v| self.depth[self.heavy[v]];
93        let (mut lo, mut hi) = if dh(u) > dh(v) { (u, v) } else { (v, u) };
94
95        while self.heavy[lo] != self.heavy[hi] {
96            lo = self.par[self.heavy[lo]];
97            if dh(lo) < dh(hi) {
98                std::mem::swap(&mut lo, &mut hi);
99            }
100        }
101        if self.depth[lo] > self.depth[hi] { hi } else { lo }
102    }
103
104    fn dfs_size(g: &mut [Vec<usize>], v: usize, size: &mut [usize]) {
105        size[v] = 1;
106        if g[v].is_empty() {
107            return;
108        }
109        for i in 0..g[v].len() {
110            let nv = g[v][i];
111            Self::dfs_size(g, nv, size);
112            size[v] += size[nv];
113            if size[nv] > size[g[v][0]] {
114                g[v].swap(0, i);
115            }
116        }
117    }
118
119    fn dfs_order(g: &[Vec<usize>], r: usize) -> Vec<usize> {
120        fn dfs(g: &[Vec<usize>], v: usize, t: &mut usize, order: &mut [usize]) {
121            order[v] = *t;
122            *t += 1;
123            for &nv in &g[v] {
124                dfs(g, nv, t, order);
125            }
126        }
127
128        let n = g.len();
129        let mut t = 0;
130        let mut order = vec![0; n];
131        dfs(g, r, &mut t, &mut order);
132        order
133    }
134
135    fn inv(p: &[usize]) -> Vec<usize> {
136        let n = p.len();
137        let mut q = vec![0; n];
138        for i in 0..n {
139            q[p[i]] = i;
140        }
141        q
142    }
143
144    fn relabel(g: &[Vec<usize>], p: &[usize]) -> Vec<Vec<usize>> {
145        let n = g.len();
146        let mut h = vec![vec![]; n];
147        for v in 0..n {
148            for &nv in &g[v] {
149                h[p[v]].push(p[nv]);
150            }
151        }
152        h
153    }
154
155    fn dfs_hld(
156        g: &[Vec<usize>],
157        r: usize,
158    ) -> (Vec<(usize, usize)>, Vec<usize>) {
159        fn dfs(
160            g: &[Vec<usize>],
161            v: usize,
162            inout: &mut [(usize, usize)],
163            next: &mut [usize],
164            t: &mut usize,
165        ) {
166            inout[v].0 = *t;
167            *t += 1;
168            for &nv in &g[v] {
169                next[nv] = if nv == g[v][0] { next[v] } else { nv };
170                dfs(g, nv, inout, next, t);
171            }
172            inout[v].1 = *t;
173            *t += 1;
174        }
175
176        let n = g.len();
177        let mut inout = vec![(0, 0); n];
178        let mut next = vec![0; n];
179        let mut t = 0;
180        dfs(g, r, &mut inout, &mut next, &mut t);
181        (inout, next)
182    }
183
184    fn par_depth(g: &[Vec<usize>], r: usize) -> (Vec<usize>, Vec<usize>) {
185        fn dfs(
186            g: &[Vec<usize>],
187            v: usize,
188            par: &mut [usize],
189            depth: &mut [usize],
190        ) {
191            for &nv in &g[v] {
192                par[nv] = v;
193                depth[nv] = depth[v] + 1;
194                dfs(g, nv, par, depth);
195            }
196        }
197
198        let n = g.len();
199        let mut par = vec![0; n];
200        let mut depth = vec![0; n];
201        dfs(g, r, &mut par, &mut depth);
202        (par, depth)
203    }
204
205    fn ascend(
206        &self,
207        mut v: usize,
208        r: usize,
209    ) -> impl Iterator<Item = HlEdge> + DoubleEndedIterator {
210        let mut res = vec![];
211        while self.heavy[v] != self.heavy[r] {
212            if self.heavy[v] == v {
213                res.push(Light(v, self.par[v]));
214                v = self.par[v];
215            } else {
216                res.push(Heavy(v, self.heavy[v]));
217                v = self.heavy[v];
218            }
219        }
220        if v != r {
221            res.push(Heavy(v, r));
222        }
223        res.into_iter()
224    }
225}
nekolib/graph/hld.rs

nekolib/graph/
hld.rs
