重心分解による分割統治法の一般形について

追記 2018/01/22

・重心を求めるアルゴリズムを変更して高速化しました。
・最後に実践的な例題を置きました。

はじめに

アリ本に載っているものよりも、（個人的に）わかりやすく整理した重心分解による分割統治法の一般形について書きます（アリ本でばっちり理解してるぜっって方は読まなくてもいい記事です）。

ちなみにこの記事は、基本的に以下の記事の実装をベースにしています。

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

重心列挙

まずは重心分解するにあたって、重心を見つけたいので、重心列挙のアルゴリズムを思い出します。それは以下のように実装できるのでした。

重心列挙の実装

vector<int> Centroid(const vector<vector<int>> &g) {
        int n = (int) 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;
}

改造

重心分解では、重心を取り除いて木を切って、その部分木を取り出していく必要があります。そこで、このアルゴリズムに少し手を加えて、任意の頂点で木を切断することができるようにします。

実際に切断するわけではなく、各頂点にdeadであるかどうかの状態を持たせて、まだ残っている頂点なのかどうかを区別することにします。その実装例は以下です。

禁止頂点を考慮した重心列挙の実装例

vector<int> Centroid(int root, const vector<vector<int>> &g, const vector<bool> &dead) {
        static vector<int> sz(g.size());
        function<void (int, int)> get_sz = [&](int u, int prev) {
                sz[u] = 1;
                for (auto v : g[u]) if (v != prev && !dead[v]) {
                        get_sz(v, u);
                        sz[u] += sz[v];
                }
        };
        get_sz(root, -1);
        int n = sz[root];
        vector<int> centroid;
        function<void (int, int)> dfs = [&](int u, int prev) {
                bool is_centroid = true;
                for (auto v : g[u]) if (v != prev && !dead[v]) {
                        dfs(v, u);
                        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(root, -1);
        return centroid;
}

まず前半で、各部分木のサイズを求めてしまっています。これでその部分木のサイズがわかるので、そのサイズを元にその部分木における重心を求めています。

高速化

ところで、上のアルゴリズムでは、重心が二つ存在する場合はその二つともを求めています。さらに、毎回$DFS$ を完全に終えてから重心を返しているため、かなりおそい実装になっています。実際には重心のうちの一つを求めればよいので、それを見つけるアルゴリズムは、以下のように実装できます。

重心を一つ求める実装例

int OneCentroid(int root, const vector<vector<int>> &g, const vector<bool> &dead) {
        static vector<int> sz(g.size()); //caution
        function<void (int, int)> get_sz = [&](int u, int prev) {
                sz[u] = 1;
                for (auto v : g[u]) if (v != prev && !dead[v]) {
                        get_sz(v, u);
                        sz[u] += sz[v];
                }
        };
        get_sz(root, -1);
        int n = sz[root];
        function<int (int, int)> dfs = [&](int u, int prev) {
                for (auto v : g[u]) if (v != prev && !dead[v]) {
                        if (sz[v] > n / 2) {
                                return dfs(v, u);
                        }
                }
                return u;
        };
        return dfs(root, -1);
}

本題

これで準備ができたので、重心分解による分割統治法を実装します。これは一般に次のように書くことができます。

重心分解による分割統治法

void CentroidDecomposition(const vector<vector<int>> &g) {
        int n = (int) g.size();
        vector<bool> dead(n, false);
        function<void (int)> rec = [&](int start) {
                int c = OneCentroid(start, g, dead);
                dead[c] = true;
                //A. compute something within each subtree alone (without the centroid)
                for (auto u : g[c]) if (!dead[u]) {
                        rec(u);
                }
                //B. compute something with the centroid
                //   piyo
                dead[c] = false;
        };
        rec(0);
}

やっていることはすごく単純で、以下のようになっています。

1. まず木全体を呼び出す。すなわち任意の頂点から再帰関数を呼び出す
2. 呼び出した木に関して、その木の重心を求めてそこで木を切る
3. 切ったときにできる各部分木内で完結する計算をする。これは各部分木を呼び出して $2$ に戻ることを意味する($A$)
4. 重心と $1$ 個以上の部分木を使ってマージする感じで計算をする($B$)
5. $DFS$ をしているので、切った部分は復元しておく

というわけで、アリ本に書かれている例題をこれを使って実装します。今回は実装がわかりやすいように、次の問題を考えることにします。

例題$1$

各辺の長さが $1$ の木が与えられる。この木の上で、長さ $k$ 以下のパスの数を数えよ

上の実装と、アリ本を参考にすると以下のような実装ができます。

int count_pairs(vector<int> &distances, int k) {
        int cnt = 0;
        sort(distances.begin(), distances.end());
        int last = (int) distances.size();
        for (int i = 0; i < (int) distances.size(); i ++) {
                while (last > 0 && distances[i] + distances[last - 1] > k) last --;
                int self = last > i;
                cnt += last - self;
        }
        return cnt / 2;
}

int CountTreePath(const vector<vector<int>> &g, int k) {
        int res = 0;
        int n = (int) g.size();
        vector<bool> dead(n, false);
        function<void (int)> rec = [&](int start) {
                int c = OneCentroid(start, g, dead);
                dead[c] = true;

                //compute something within a subtree alone (without the centroid)
                for (auto u : g[c]) if (!dead[u]) {
                        rec(u);
                }

                //compute something with the centroid
                vector<int> distances;
                distances.push_back(0);
                for (auto u : g[c]) if (!dead[u]) {
                        vector<int> tmp_distances;
                        function<void (int, int, int)> enumerate_paths = [&](int u, int prev, int distance) {
                                tmp_distances.push_back(distance);
                                for (auto v : g[u]) if (v != prev && !dead[v]) {
                                        enumerate_paths(v, u, distance + 1);
                                }
                        };
                        enumerate_paths(u, c, 1);
                        res -= count_pairs(tmp_distances, k);
                        distances.insert(distances.end(), tmp_distances.begin(), tmp_distances.end());
                }
                res += count_pairs(distances, k);
                //

                dead[c] = false;
        };
        rec(0);
        return res;
}

結局ほとんど $B$ の部分しか書き換える必要がないことがわかりますね。

以上で重心分解による分割統治法が簡単にできました。

例題$2$

さて、以下の問題は重心分解で解くことができます。

CS Academy

さっきは長さ $k$ 以下のパスの個数を求めましたが、今度は長さがちょうど $k$ であるようなパスの個数を求めなければいけません。

ところが長さ $k - 1$ 以下であるようなパスの個数もさっきのアルゴリズムで求めることができるので、その個数を引けば良いだけです。全体の実装は以下のようになります。

#include <cstdio>
#include <algorithm>
#include <vector>
#include <functional>
using namespace std;

int OneCentroid(int root, const vector<vector<int>> &g, const vector<bool> &dead) {
        static vector<int> sz(g.size()); //caution
        function<void (int, int)> get_sz = [&](int u, int prev) {
                sz[u] = 1;
                for (auto v : g[u]) if (v != prev && !dead[v]) {
                        get_sz(v, u);
                        sz[u] += sz[v];
                }
        };
        get_sz(root, -1);
        int n = sz[root];
        function<int (int, int)> dfs = [&](int u, int prev) {
                for (auto v : g[u]) if (v != prev && !dead[v]) {
                        if (sz[v] > n / 2) {
                                return dfs(v, u);
                        }
                }
                return u;
        };
        return dfs(root, -1);
}

int count_pairs(vector<int> &distances, int k) {
        int cnt = 0;
        sort(distances.begin(), distances.end());
        int last = (int) distances.size();
        for (int i = 0; i < (int) distances.size(); i ++) {
                while (last > 0 && distances[i] + distances[last - 1] > k) last --;
                int self = last > i;
                cnt += last - self;
        }
        return cnt / 2;
}


int CountTreePath(const vector<vector<int>> &g, int k) {
        int res = 0;
        int n = (int) g.size();
        vector<bool> dead(n, false);
        function<void (int)> rec = [&](int start) {
                int c = OneCentroid(start, g, dead);
                dead[c] = true;
                //compute something within a subtree alone (without the centroid)
                for (auto u : g[c]) if (!dead[u]) {
                        rec(u);
                }
                //compute something with the centroid
                vector<int> distances;
                distances.push_back(0);
                for (auto u : g[c]) if (!dead[u]) {
                        vector<int> tmp_distances;
                        function<void (int, int, int)> enumerate_paths = [&](int u, int prev, int distance) {
                                tmp_distances.push_back(distance);
                                for (auto v : g[u]) if (v != prev && !dead[v]) {
                                        enumerate_paths(v, u, distance + 1);
                                }
                        };
                        enumerate_paths(u, c, 1);
                        res -= count_pairs(tmp_distances, k);
                        res += count_pairs(tmp_distances, k - 1);
                        distances.insert(distances.end(), tmp_distances.begin(), tmp_distances.end());
                }
                res += count_pairs(distances, k);
                res -= count_pairs(distances, k - 1);
                //
                dead[c] = false;
        };
        rec(0);
        return res;
}

const int MOD = 1e9 + 7;

int main() {
        int n, k;
        scanf(&quot;%d%d&quot;, &n, &k);
        vector<vector<int>> g(n);
        for (int i = 0; i < n - 1; i ++) {
                int a, b;
                scanf(&quot;%d%d&quot;, &a, &b);
                a --, b --;
                g[a].push_back(b);
                g[b].push_back(a);
        }
        int cnt = CountTreePath(g, k);
        long long ans = ((long long) k * (k + 1) / 2) % MOD * cnt % MOD;
        printf(&quot;%lld\n&quot;, ans);
        return 0;
}

最後に練習問題を置いておくので、これを使ってみたいという方がいれば、使ってみてください（ただしふつうにアルゴリズムが難しい問題です）。

Problem - E - Codeforces

解説は以下に書いていますので、参考までに。

Codeforces 458 E. Palindromes in a Tree - Learning Algorithms

以上です。