pub mod data_structure {
#[derive(Debug)]
pub struct HeavyLightDecomposition {
pre_order: Vec<usize>, // pre_order[i] = 頂点 i の訪問時刻
parent: Vec<Option<usize>>, // parent[i] = 頂点 i の親の頂点番号
head: Vec<usize>, // head[i] = 頂点 i と同じ heavy path における最も根側の頂点番号
pub num_childs: Vec<usize>, // num_childs[i] = 頂点 i を根とする部分木の頂点数
edges: Vec<Vec<usize>>,
}
impl HeavyLightDecomposition {
pub fn new(n: usize) -> Self {
assert!(n > 0);
Self {
pre_order: vec![usize::MAX; n],
parent: vec![None; n],
head: vec![usize::MAX; n],
num_childs: vec![1; n],
edges: vec![vec![]; n],
}
}
/// u -> v と v -> u 方向に辺を追加する
pub fn add_edge(&mut self, u: usize, v: usize) {
self.edges[u].push(v);
self.edges[v].push(u);
}
pub fn build(&mut self) {
self.dfs_for_num_childs(0, None);
let mut order = 0;
self.dfs(0, None, 0, &mut order);
}
/// 頂点 v の訪問時刻を返す
pub fn index(&self, v: usize) -> usize {
self.pre_order[v]
}
/// 頂点 u と頂点 v の最小共通祖先を返す
pub fn lca(&self, mut u: usize, mut v: usize) -> usize {
loop {
if self.pre_order[u] > self.pre_order[v] {
std::mem::swap(&mut u, &mut v);
}
if self.head[u] == self.head[v] {
return u;
}
v = self.parent[self.head[v]].unwrap();
}
}
/// 頂点 from から頂点 to へのパスをいくつかの heavy path に分解し、各 heavy path の端点を返すイテレータを返す。
pub fn iter_heavy_path(&self, from: usize, to: usize) -> PathSegmentsIterator {
PathSegmentsIterator {
hld: self,
from,
to,
done: false,
}
}
/// num_childs を求めるための dfs
fn dfs_for_num_childs(&mut self, v: usize, parent: Option<usize>) {
let n = self.edges[v].len();
for i in 0..n {
let to = self.edges[v][i];
if Some(to) == parent {
continue;
}
self.dfs_for_num_childs(to, Some(v));
self.num_childs[v] += self.num_childs[to];
}
}
/// pre_order, parent, head を求めるための dfs
/// Self::dfs_for_num_childs の後に実行する
fn dfs(&mut self, v: usize, parent: Option<usize>, head: usize, order: &mut usize) {
self.parent[v] = parent;
self.pre_order[v] = *order;
*order += 1;
self.head[v] = head;
let next_id = {
let mut max = 0;
let mut max_id = None;
for &to in &self.edges[v] {
if Some(to) == parent {
continue;
}
if max < self.num_childs[to] {
max = self.num_childs[to];
max_id = Some(to);
}
}
max_id
};
if let Some(to) = next_id {
self.dfs(to, Some(v), head, order);
}
let n = self.edges[v].len();
for i in 0..n {
let to = self.edges[v][i];
if Some(to) == next_id || Some(to) == parent {
continue;
}
self.dfs(to, Some(v), to, order);
}
}
}
pub struct PathSegmentsIterator<'a> {
hld: &'a HeavyLightDecomposition,
from: usize,
to: usize,
done: bool,
}
impl Iterator for PathSegmentsIterator<'_> {
/// (根側の頂点番号, 葉側の頂点番号 (端点を含む), lca を含む path か, 元のパスの u -> v の向きと同じか)
type Item = (usize, usize, bool, bool);
fn next(&mut self) -> Option<Self::Item> {
let Self {
hld,
from,
to,
done,
} = *self;
(!done).then_some(())?;
let HeavyLightDecomposition {
pre_order,
parent,
head,
..
} = hld;
if head[from] == head[to] {
self.done = true;
if pre_order[from] < pre_order[to] {
Some((from, to, true, false))
} else {
Some((to, from, true, true))
}
} else if pre_order[from] < pre_order[to] {
self.to = parent[head[to]].unwrap();
Some((head[to], to, false, false))
} else {
self.from = parent[head[from]].unwrap();
Some((head[from], from, false, true))
}
}
}
}