概要
本記事では、木上のすべてのパスについての計算を効率的に行う手法の一つである1/3重心分解の実装例について紹介します。
アルゴリズムの詳細や計算量解析はこの記事では省略します。先に参考記事をご覧ください。
重心分解とは
重心分解は、木の全てのパスに関する計算を、より扱いやすい「根付き木のすべての根を通るパスに関する計算」へ分解するアルゴリズムです。
具体的には以下の手続きをします。
- 重心を求める
- 重心を取り除く
- できた各部分木に対して、再帰的に重心分解を実行する
「重心を取り除く」操作は、「頂点を無視するフラグ」を立てることで簡単に実現できます。
1/3重心分解とは
一方1/3重心分解は、「木上のすべてのパスについての計算」を「赤色の頂点と青色の頂点を結ぶパスについての計算」に分解するアルゴリズムです。
具体的には以下の手続きをします。
- 重心の周りの部分木を、辺の数がだいたい半分ずつに分かれるように赤と青に塗り分ける
- 各色の頂点と根を合わせた部分木について、再帰的に1/3重心分解を実行する
1/3重心分解の問題点
1/3重心分解は重心分解とは異なり、重心として選ばれた頂点は再び分解後の2つの部分木に現れます。
そのため重心分解のような「頂点を無視するフラグ」を用いた実装では計算量が悪化してしまいます。
これの対処として、陽に部分木を構築する方法がよく用いられているようです。
しかし、陽に部分木を構築するための時間やメモリがかかったり、頂点番号の振りなおしなど煩雑になりがちな実装が必要になったりします。
以下では木を工夫して持つことで、陽に部分木を構築することなく、頂点番号の振りなおしも不要にする方法を紹介します。
木の持ち方
部分木を表現するために、CSR形式 (Compressed Sparse Row) で木を管理します。
CSR形式はグラフを以下のような1次元配列の組で表現する形式です。
-
dest: 辺を始点の昇順にsortしたときの終点を並べたもの -
head[v]: 頂点vの最初の辺のインデックス -
tail[v]: 頂点vの最後の辺のインデックス
頂点 v から出る辺は dest[head[v]..tail[v]] に格納されています。
このように持つことで、ある頂点から出る辺のswapと有効な辺の範囲の変更がO(1)で行えます。
赤色と青色の頂点がそれぞれまとまるように辺を並び替え、次の部分木に移行するときに有効な辺の範囲を変更することで、差分更新で部分木を表現できます。
実装例
以下に1/3重心分解の実装例を示します。
※ 実装の都合上、長さ1以下のパスは列挙していないので、別途計算するなどの対処が必要です。
// 1/3重心分解
#[derive(Clone, Copy)]
struct CentroidDecomposition<'a> {
root: usize,
vertices: [&'a [usize]; 2],
parent: &'a [usize],
}
impl<'a> CentroidDecomposition<'a> {
fn run(edges: &[(usize, usize)], mut f: impl FnMut(CentroidDecomposition)) {
Tree::new(edges).dfs_decompose(0, &mut f);
}
}
struct Tree {
dest: Vec<usize>,
head: Vec<usize>,
tail: Vec<usize>,
size: Vec<usize>,
parent: Vec<usize>,
}
impl Tree {
fn new(edges: &[(usize, usize)]) -> Self {
let n = edges.len() + 1;
let mut tail = vec![0; n];
for &(u, v) in edges {
tail[u] += 1;
tail[v] += 1;
}
for i in 1..n {
tail[i] += tail[i - 1];
}
let mut head = tail.clone();
let mut dest = vec![0; (n - 1) * 2];
for &(u, v) in edges.iter().rev() {
head[u] -= 1;
dest[head[u]] = v;
head[v] -= 1;
dest[head[v]] = u;
}
Self {
dest,
head,
tail,
size: vec![0; n],
parent: vec![!0; n],
}
}
fn dfs_decompose(&mut self, v: usize, f: &mut impl FnMut(CentroidDecomposition)) {
// 今見ている木を v を根としたときの部分木サイズと親を求める
self.parent[v] = !0;
let mut ord = vec![v];
let mut i = 0;
while i < ord.len() {
let v = ord[i];
self.size[v] = 1;
for &u in &self.dest[self.head[v]..self.tail[v]] {
if u != self.parent[v] {
self.parent[u] = v;
ord.push(u);
}
}
i += 1;
}
let n = ord.len();
if n <= 2 {
return; // 重心が葉の場合は終了
}
// 重心を求める
let mut c = !0;
for &v in ord.iter().rev() {
if self.parent[v] != !0 {
self.size[self.parent[v]] += self.size[v];
}
if c == !0 && self.size[v] >= (n + 1) / 2 {
c = v;
}
}
// 根を重心に変更する
let mut p = !0;
let mut p_size = 0;
let mut v = c;
while v != !0 {
std::mem::swap(&mut self.size[v], &mut p_size);
self.size[v] = n - self.size[v];
std::mem::swap(&mut self.parent[v], &mut p);
std::mem::swap(&mut p, &mut v);
}
// 重心を根とした木 2 つに分解する
// 重心から出る辺を前半と後半に分ける
let mut sum_size = 0;
let mut mid = self.head[c];
for i in self.head[c]..self.tail[c] {
let v = self.dest[i];
if sum_size + self.size[v] <= (n - 1) / 2 {
sum_size += self.size[v];
self.dest.swap(mid, i);
mid += 1;
}
}
// 各部分木の頂点をBFS順に並べる
let vertices =
[&self.dest[self.head[c]..mid], &self.dest[mid..self.tail[c]]].map(|children| {
let mut vertices = children.to_vec();
let mut i = 0;
while i < vertices.len() {
let v = vertices[i];
for &u in &self.dest[self.head[v]..self.tail[v]] {
if u != self.parent[v] {
vertices.push(u);
}
}
i += 1;
}
vertices
});
// 赤の頂点から青の頂点のパスすべてについて計算する
f(CentroidDecomposition {
root: c,
vertices: [&vertices[0], &vertices[1]],
parent: &self.parent,
});
// 各部分木について再帰的に実行する
std::mem::swap(&mut self.head[c], &mut mid);
self.dfs_decompose(c, f);
std::mem::swap(&mut self.head[c], &mut mid);
std::mem::swap(&mut self.tail[c], &mut mid);
self.dfs_decompose(c, f);
std::mem::swap(&mut self.tail[c], &mut mid);
}
}
使用例
ABC359 G - Sum of Tree Distance
まとめ
1/3重心分解のCSR形式を用いた実装例を紹介しました。
比較的シンプルな実装になっていると思います。
参考になりましたら幸いです。