木の重心列挙アルゴリズム

木の重心列挙アルゴリズム

English version is available here.

木の重心を列挙するアルゴリズムです。重心の性質をよく知っている方にとっては自明な木$dp$を実装するだけだと思います。

これで $verify$ しています。計算量は $O(n)$ です。

中のアルゴリズムとしては、基本的には各頂点についてその頂点からのびる部分木のサイズがすべて $\cfrac {n}{2}$ 以下であるような頂点を求めているだけです。

では、そのような頂点が $3$ 個以上存在しないことを示しておきます。

まず、重心が $2$ 個あるとき、それらは必ず隣接していることを示します。

仮に以下の画像のように、重心 $u$ と重心 $v$ が存在して、これらが隣接していないと仮定します。すると $u$ と $v$ のパス上には必ず一つ以上の頂点が存在することになります。そのうちの一つを $w$ とします。

$v$ の部分木であって $u$ を含むものを部分木 $V$ 、$u$ の部分木であって $v$ を含むものを部分木 $U $ と呼ぶことにすると、 $v$ , $u$ はそれぞれ重心なので、以下が言えます。

$|U| \leq \cfrac{n}{2}$
$|V| \leq \cfrac{n}{2}$
$\therefore |U| + |V| \leq n$

ところで仮定より、$W = U \cap V$ とすると、$|W| >= 1$ であることから、$|U| + |V| = |U \cup V| - |W| = n - |W| < n$ となるので、これは矛盾します。よって重心が $2$ 個存在してかつそれらが隣接していないということはありえないことが言えました。

隣接しているときに限り、$|W| = 0$ となって、条件を満たします。このとき明らかに $|V| = |U| = \cfrac {n}{2}$ となります。

ここまで言えば $3$ 個以上存在しないことは明らかだと思います。

以上の実装は、根を固定して $DFS$ しながら、各頂点以下の部分木のサイズと、親側の部分木のサイズ（これは $n$ からその部分木のサイズを引けばよいです）を順に求めていくような木$dp$で求めることができます。

実装

vector<int> Centroid(const vector<vector<int>> &g) {
        int n = g.size();
        vector<int> centroid;
        vector<int> sz(n);
        function<void (int, int)> dfs = [&](int u, int prev) {
                sz[u] = 1;
                bool is_centroid = true;
                for (auto v : g[u]) if (v != prev) {
                        dfs(v, u);
                        sz[u] += sz[v];
                        if (sz[v] > n / 2) is_centroid = false;
                }
                if (n - sz[u] > n / 2) is_centroid = false;
                if (is_centroid) centroid.push_back(u);
        };
        dfs(0, -1);
        return centroid;
}

引数のgは木の隣接リストで、それを投げればその重心が1個または2個返ってきます。

使用例ですが、上にも書いた $verify$ に使った問題の他に、次のような問題もあるので見てみることにします。

D - Tree and Hamilton Path

まず各辺について、その通過回数の上界を考えてみることにします。その辺が繋ぐ二つの部分木を $S$, $T$ とすると、できるだけ行き来することを考えると最大 $2 * min(|S|, |T|)$ 回通過することができます。もう少し考えると重心からスタートすればほとんどすべての辺でこの上界を達成でき、二つの重心間の辺または一つの重心に繋がる任意の一辺だけ通過回数が $1$ 回少なくなることがわかります。

この問題では結局それぞれの部分木のサイズを求めなければいけませんが、重心はなにも考えずに求めることができるので考えるのがかなり楽になると思います。

実装

#include <iostream>
#include <cstdio>
#include <algorithm>
#include <string>
#include <vector>
#include <map>
#include <set>
#include <functional>
#include <cassert>
using namespace std;
 
vector<int> Centroid(const vector<vector<int>> &g) {
        int n = g.size();
        vector<int> centroid;
        vector<int> sz(n);
        function<void (int, int)> dfs = [&](int u, int prev) {
                sz[u] = 1;
                bool is_centroid = true;
                for (auto v : g[u]) if (v != prev) {
                        dfs(v, u);
                        sz[u] += sz[v];
                        if (sz[v] > n / 2) is_centroid = false;
                }
                if (n - sz[u] > n / 2) is_centroid = false;
                if (is_centroid) centroid.push_back(u);
        };
        dfs(0, -1);
        return centroid;
}
 
int main() {
        int n;
        scanf("%d", &n);
        vector<vector<int>> g(n);
        vector<vector<pair<int ,int>>> gg(n);
        for (int i = 0; i < n - 1; i ++) {
                int a, b, c;
                scanf("%d%d%d", &a, &b, &c);
                a --, b --;
                g[a].push_back(b);
                g[b].push_back(a);
                gg[a].emplace_back(b, c);
                gg[b].emplace_back(a, c);
        }
        long long ans = 0;
        vector<int> sz(n);
        function<void (int, int)> dfs = [&](int u, int prev) {
                sz[u] = 1;
                for (auto e : gg[u]) {
                        int v = e.first, cost = e.second;
                        if (v == prev) continue;
                        dfs(v, u);
                        sz[u] += sz[v];
                        ans += min(sz[v], n - sz[v]) * (long long) 2 * cost;
                }                                 
        };
        dfs(0, -1);
        auto centroid = Centroid(g);
        if (centroid.size() == 1) {
                int mi = 0x3f3f3f3f;
                for (auto e : gg[centroid[0]]) {
                        mi = min(mi, e.second);
                }
                ans -= mi;
        } else if (centroid.size() == 2) {
                for (auto e : gg[centroid[0]]) {
                        if (e.first == centroid[1]) {
                                ans -= e.second;
                                break;
                        }
                }
        } else {
                assert(false);
        }
        printf("%lld\n", ans);
        return 0;
}

こちらの問題も重心を見つける必要がある問題です。

第4回ドワンゴからの挑戦状予選 E. ニワンゴくんの家探し - Learning Algorithms

動的に木を変化させなければいけないところが注意点で、重心判定のときの $\frac {n}{2}$ との比較で、一番最初の木のサイズ $n$ を使わないようにしなければいけません。