haar_lib/tree/
depth_query.rs

1//! Tree depth query
2//!
3//! # References
4//! - [https://niuez.github.io/posts/entry/2019/10/05/002503/](https://niuez.github.io/posts/entry/2019/10/05/002503/)
5//! - [https://niuez.github.io/posts/dfs_bfs_et/](https://niuez.github.io/posts/dfs_bfs_et/)
6//!
7//! # Problems
8//! - [yukicoder No.899 γatheree](https://yukicoder.me/problems/no/899)
9
10use crate::tree::*;
11use std::collections::VecDeque;
12
13/// 根付き木において、同一の深さの頂点の区間に対して区間クエリができる。
14pub struct TreeDepthQuery {
15    par: Vec<Option<usize>>,
16    depth: Vec<usize>,
17    left: Vec<usize>,
18    right: Vec<usize>,
19    bfs_ord: Vec<Vec<usize>>,
20    dfs_ord: Vec<Vec<usize>>,
21    ord: Vec<usize>,
22}
23
24impl TreeDepthQuery {
25    /// 根を`root`とする`tree`を基に、`TreeDepthQuery`を構築する。
26    pub fn new<E: TreeEdgeTrait>(tree: &Tree<E>, root: usize) -> Self {
27        let size = tree.len();
28        let mut this = Self {
29            par: vec![None; size],
30            depth: vec![0; size],
31            left: vec![0; size],
32            right: vec![0; size],
33            bfs_ord: vec![],
34            dfs_ord: vec![],
35            ord: vec![0; size],
36        };
37
38        this.dfs(tree, root, None, 0, &mut 0);
39
40        let mut q = VecDeque::new();
41        q.push_back((root, 0));
42        let mut ord = 0;
43
44        while let Some((i, d)) = q.pop_front() {
45            if this.bfs_ord.len() <= d {
46                this.bfs_ord.push(vec![]);
47            }
48            this.bfs_ord[d].push(ord);
49            this.ord[i] = ord;
50            ord += 1;
51
52            for e in tree.nodes[i].neighbors() {
53                if this.par[i].is_none_or(|p| p != e.to()) {
54                    q.push_back((e.to(), d + 1));
55                }
56            }
57        }
58
59        this
60    }
61
62    fn dfs<E: TreeEdgeTrait>(
63        &mut self,
64        tree: &Tree<E>,
65        cur: usize,
66        par: Option<usize>,
67        d: usize,
68        ord: &mut usize,
69    ) {
70        self.par[cur] = par;
71        self.depth[cur] = d;
72
73        if self.dfs_ord.len() <= d {
74            self.dfs_ord.push(vec![]);
75        }
76        self.dfs_ord[d].push(*ord);
77        self.left[cur] = *ord;
78        *ord += 1;
79
80        for e in tree.nodes[cur].neighbors() {
81            if par.is_none_or(|p| p != e.to()) {
82                self.dfs(tree, e.to(), Some(cur), d + 1, ord);
83            }
84        }
85
86        self.right[cur] = *ord;
87    }
88
89    /// 頂点`i`から深さ`d`の子孫頂点に対応する区間を返す。
90    ///
91    /// **Time complexity** $O(\log n)$
92    pub fn children_query(&self, i: usize, d: usize) -> Option<(usize, usize)> {
93        let d = d + self.depth[i];
94        if self.bfs_ord.len() > d {
95            let l = match self.dfs_ord[d].binary_search(&self.left[i]) {
96                Ok(x) | Err(x) => x,
97            };
98            let r = match self.dfs_ord[d].binary_search(&self.right[i]) {
99                Ok(x) | Err(x) => x,
100            };
101
102            if l >= self.bfs_ord[d].len() {
103                return None;
104            }
105            if r == 0 {
106                return None;
107            }
108
109            Some((self.bfs_ord[d][l], self.bfs_ord[d][r - 1] + 1))
110        } else {
111            None
112        }
113    }
114
115    /// 頂点`i`に対応する区間を返す。
116    pub fn me_query(&self, i: usize) -> (usize, usize) {
117        (self.ord[i], self.ord[i] + 1)
118    }
119
120    /// 頂点`i`の`k`個遡った祖先の頂点を返す。
121    ///
122    /// **Time complexity** $O(k)$
123    pub fn ancestor(&self, i: usize, k: usize) -> Option<usize> {
124        let mut p = i;
125        for _ in 0..k {
126            match self.par[p] {
127                Some(x) => p = x,
128                _ => return None,
129            }
130        }
131        Some(p)
132    }
133}
haar_lib/tree/depth_query.rs

haar_lib/tree/
depth_query.rs
