haar_lib/tree/
hld.rs

1//! 重軽分解
2use crate::tree::*;
3use std::cmp::max;
4
5/// 重軽分解
6#[derive(Clone, Debug)]
7pub struct HLD {
8    _size: usize,
9    par: Vec<Option<usize>>,
10    head: Vec<usize>,
11    id: Vec<usize>,
12    rid: Vec<usize>,
13    next: Vec<Option<usize>>,
14    end: Vec<usize>,
15}
16
17impl HLD {
18    /// **Time complexity** $O(n)$
19    ///
20    /// **Space complexity** $O(n)$
21    pub fn new<E: TreeEdgeTrait>(tree: &Tree<E>, root: usize) -> Self {
22        let size = tree.len();
23        let mut ret = Self {
24            _size: size,
25            par: vec![None; size],
26            head: vec![0; size],
27            id: vec![0; size],
28            rid: vec![0; size],
29            next: vec![None; size],
30            end: vec![0; size],
31        };
32
33        let mut tr = vec![vec![]; size];
34        for (i, nodes) in tree.nodes.iter().enumerate() {
35            for e in nodes.neighbors() {
36                tr[i].push(e.to());
37            }
38        }
39
40        ret.dfs_sub(&mut tr, root, None, &mut vec![1; size]);
41        ret.dfs_build(&tr, root, &mut 0);
42        ret
43    }
44
45    fn dfs_sub(
46        &mut self,
47        tree: &mut [Vec<usize>],
48        cur: usize,
49        par: Option<usize>,
50        sub: &mut Vec<usize>,
51    ) {
52        self.par[cur] = par;
53        tree[cur].retain(|&x| Some(x) != par);
54
55        let mut t = 0;
56        let n = tree[cur].len();
57        for i in 0..n {
58            let to = tree[cur][i];
59            self.dfs_sub(tree, to, Some(cur), sub);
60            sub[cur] += sub[to];
61            if sub[to] > t {
62                t = sub[to];
63                self.next[cur] = Some(to);
64                tree[cur].swap(i, 0);
65            }
66        }
67    }
68
69    fn dfs_build(&mut self, tree: &[Vec<usize>], cur: usize, index: &mut usize) {
70        self.id[cur] = *index;
71        self.rid[*index] = cur;
72        *index += 1;
73
74        for (i, &to) in tree[cur].iter().enumerate() {
75            self.head[to] = if i == 0 { self.head[cur] } else { to };
76            self.dfs_build(tree, to, index);
77        }
78
79        self.end[cur] = *index;
80    }
81
82    /// 演算は可換性を仮定する。
83    ///
84    /// **Time complexity** $O(\log n)$
85    pub fn path_query_vertex(&self, mut x: usize, mut y: usize) -> Vec<(usize, usize)> {
86        let mut ret = vec![];
87        loop {
88            if self.id[x] > self.id[y] {
89                std::mem::swap(&mut x, &mut y);
90            }
91            ret.push((max(self.id[self.head[y]], self.id[x]), self.id[y] + 1));
92            if self.head[x] == self.head[y] {
93                break;
94            }
95            y = self.par[self.head[y]].unwrap();
96        }
97        ret
98    }
99
100    /// **Time complexity** $O(\log n)$
101    pub fn path_query_edge(&self, mut x: usize, mut y: usize) -> Vec<(usize, usize)> {
102        let mut ret = vec![];
103        loop {
104            if self.id[x] > self.id[y] {
105                std::mem::swap(&mut x, &mut y);
106            }
107            if self.head[x] == self.head[y] {
108                if x != y {
109                    ret.push((self.id[x] + 1, self.id[y] + 1));
110                }
111                break;
112            }
113            ret.push((self.id[self.head[y]], self.id[y] + 1));
114            y = self.par[self.head[y]].unwrap();
115        }
116        ret
117    }
118
119    /// **Time complexity** $O(1)$
120    pub fn subtree_query_vertex(&self, x: usize) -> (usize, usize) {
121        (self.id[x], self.end[x])
122    }
123
124    /// **Time complexity** $O(1)$
125    pub fn subtree_query_edge(&self, x: usize) -> (usize, usize) {
126        (self.id[x] + 1, self.end[x])
127    }
128
129    /// **Time complexity** $O(1)$
130    pub fn parent(&self, x: usize) -> Option<usize> {
131        self.par[x]
132    }
133
134    /// **Time complexity** $O(1)$
135    pub fn get_id(&self, x: usize) -> usize {
136        self.id[x]
137    }
138
139    /// **Time complexity** $O(1)$
140    pub fn get_edge_id(&self, u: usize, v: usize) -> Option<usize> {
141        if self.par[u] == Some(v) {
142            Some(self.id[u])
143        } else if self.par[v] == Some(u) {
144            Some(self.id[v])
145        } else {
146            None
147        }
148    }
149
150    /// **Time complexity** $O(\log n)$
151    pub fn lca(&self, mut u: usize, mut v: usize) -> usize {
152        loop {
153            if self.id[u] > self.id[v] {
154                std::mem::swap(&mut u, &mut v);
155            }
156            if self.head[u] == self.head[v] {
157                return u;
158            }
159            v = self.par[self.head[v]].unwrap();
160        }
161    }
162}
haar_lib/tree/hld.rs

haar_lib/tree/
hld.rs
