ARC101 E問題 : Ribbons on Tree - 忘れても大丈夫

問題. Ribbons on Tree

$n$ 頂点の木が与えられる．木のどの頂点もちょうど1つの頂点対に割り当てられるような $N / 2$ 個の頂点対への分割を考える．そのような分割で，どの辺もある組の最短経路上にあるような分割が何通りあるかを求めよ．

制約： $2 \le n \le 5,000$

解法1. 包除原理

包除原理までは思いついたのですが動的計画法で解く方法が思いつかなかったので解説を見ました [1]．解説のマージの部分が分からなかったので kmjp さんのソースコードを参考にしました．
　＊ kmjpさんの記事がありました：AtCoder ARC #101 : E - Ribbons on Tree - kmjp's blog

木 $T = (V, E)$ を考えます．辺 $e \in E$ に対して， $A_{e}$ を最短経路が $e$ を通らない頂点対への分割全体とします．このとき， $e$ 以外の辺に関してはその辺を通る頂点対の最短経路が存在するかどうかは考えていません．頂点数が偶数 $n'$ の頂点集合に対して， $g(n')$ を頂点対への分割の組合せ数とすると，
　 $g(n') := (n' - 1) \times (n' - 3) \times \cdots \times 1$
となります． $n'$ が奇数の場合は $g(n') = 0$ で， $\left| A_{\emptyset} \right|$ は $g(n)$ と等しくなります．
$A_{e}$ は $e$ を通る最短経路が存在しない頂点対への分割なので， $T$ から $e$ を除去したグラフ上の分割の仕方になります．一般に木から $F \subseteq E$ を取り除くと， $|F| + 1$ 個の木に分割されるので，各連結成分の頂点数を $n_1, n_2, \ldots, n_{|F|+1}$ とすると，
　 $\left| \bigcap_{e \in F} A_e \right| := g(n_1) g(n_2) \cdots g(n_{|F|+1})$
となります．したがって，包除原理から求める答えは，
　 $\left| \bigcap_{e \in E} \overline{A_e} \right| = \left| \overline{\bigcap_{e \in E} A_e} \right| = \sum_{F \subseteq E} (-1)^{|F|} \left| \bigcap_{e \in F} A_e \right|$
となります．包除原理に従って実装すると $O(n 2^n)$ 時間かかります．

計算時間： $O(n 2^n)$

ソースコードを表示 : 包除原理

#include <bits/stdc++.h>

using namespace std;

using ll = long long;
constexpr ll MOD = 1e9 + 7;

class ProblemE{
public:
    ProblemE(int _n) : n(_n), m(n - 1), tree(m), comb(n + 1), visited(n) {
        for (int i = 0; i < m; ++i) {
            cin >> tree[i].first >> tree[i].second;
            --tree[i].first; --tree[i].second;
        }

        comb[2] = 1;
        for (int i = 4; i <= n; i += 2)
            comb[i] = ((i - 1) * comb[i - 2]) % MOD;
    }

    ll Solve() {
        ll ans = 0;
        for (int s = 0; s < 1 << m; ++s) {
            ll sum = 1;
            fill(visited.begin(), visited.end(), false);

            for (int v = 0; v < n; ++v) {
                if (visited[v]) continue;
                int num_visited = Bfs(v, s);
                sum = (sum * comb[num_visited]) % MOD;
            }

            if (__builtin_popcount(s) % 2 != 0) ans = ((ans - sum) + MOD) % MOD;
            else ans = (ans + sum) % MOD;
        }
        return ans;
    }

private:
    int n, m;
    vector<pair<int, int>> tree;
    vector<int> comb;
    vector<bool> visited;

    int Bfs(const int cur, const int ban) {
        visited[cur] = true;
        int res = 1;

        for (int i = 0; i < m; ++i) {
            int dst = (tree[i].first == cur) ? tree[i].second : -1;
            if (tree[i].second == cur) dst = tree[i].first;

            if (dst == -1 || visited[dst] || ban >> i & 1) continue;
            res += Bfs(dst, ban);
        }
        return res;
    }

};


int main() {
    cin.tie(0); ios::sync_with_stdio(false);

    int n; cin >> n;
    assert(n <= 20);

    ProblemE p(n);
    cout << p.Solve() << endl;

    return 0;
}

解法2. 木上の動的計画法

頂点 $r \in V$ を根とする根付き木 $T$ を考えます．また，頂点 $v$ を根とする部分木のサイズを $s(v)$ とします． $T$ の各部分木に対して上の包除原理を適用すること考えるのですが， $\sum_{F \subseteq E} (-1)^{|F|} \left| \bigcap_{e \in F} A_e \right|$ から根と子をマージするときに必要な情報は，根を含む連結成分のサイズと，取り除いた辺数のパリティであることが分かります．
$f(v, x, p)$ を頂点 $v$ を根とする部分木で， $v$ を含む連結成分の頂点数が $x$ で，その連結成分以外に対して，どの頂点対の最短経路にも含まれない辺数のパリティが $p$ となる分割の組合せ数とします．この漸化式が計算できると答えは， $\sum_{x=0}^{n} (f(r, x, 0) - f(r, x, 1)) \times g(x)$ となります．
頂点 $v$ に対して $f(v, \cdot, \cdot)$ を計算します． $v$ の子全体を $c_1 < c_2 < \cdots < c_k$ とします（ $<$ は $v$ の子全体に対しての適当な全順序）． $1 \le i < k$ に対して， $c_1, \ldots, c_i$ の根付き木と頂点 $v$ を併せた部分木 $T_{v, i}$ に対して $f(v, 0 \ldots (s(c_1) + \cdots + s(c_i) + 1), \cdot)$ までの値が求まっているとして，子 $c_{i+1}$ の根付き木 $T_{c_{i+1}}$ をその部分木にマージすることを考えます．これは， $T_{v, i}$ の $v$ を含むサイズ $1 \le x \le s(c_1) + \cdots + s(c_i) + 1$ の連結成分と， $T_{c_{i+1}}$ の $c_{i+1}$ をサイズ $1 \le y \le s(c_{i+1})$ の連結成分に対して次の2通りのマージの仕方があります．

1. 辺 $\{v, c_{i+1}\} \in E$ を使用
　・ $f(v, x + y, 0) += f(v, x, 0) \times f(c_{i+1}, y, 0) + f(v, x, 1) \times f(c_{i+1}, y, 1)$
　・ $f(v, x + y, 1) += f(v, x, 0) \times f(c_{i+1}, y, 1) + f(v, x, 1) \times f(c_{i+1}, y, 0)$

2. 辺 $\{v, c_{i+1}\} \in E$ を不使用
　・ $f(v, x, 1) += (f(v, x, 0) \times f(c_{i+1}, y, 0) + f(v, x, 1) \times f(c_{i+1}, y, 1)) \times g(y)$
　・ $f(v, x, 0) += (f(v, x, 0) \times f(c_{i+1}, y, 1) + f(v, x, 1) \times f(c_{i+1}, y, 0)) \times g(y)$

2番目のマージで $g(y)$ の項を掛けているのは，辺 $\{v, c_{i+1}\}$ を使用しないので $T_{c_{i+1}}$ の $c_{i+1}$ を含むサイズ $y$ の連結成分の頂点対の分割の仕方が定まるためです．
上の漸化式を深さ優先探索で後順で更新すると計算できます．計算時間は一見すると各頂点で $O(n^2)$ ステップかかり，全体で $O(n^3)$ かかりそうですが解析をすると $O(n^2)$ 時間となります．競技プログラミング界隈では「木 DP 二乗」として有名らしいです（二乗の木 DP - (iwi) ｛反省します - TopCoder部）．
簡単のため分枝数が2で，左部分木と右部分木のサイズがそれぞれ $l, r$ として計算時間の解析を行います． $T(n')$ を部分木のサイズが $n'$ のときの最悪実行時間とすると，上の漸化式から $T(l + r) = T(l) + T(r) + l \times r$ となります．ある正整数 $c$ に対して， $T(n') \le cn^2$ を帰納法で示します．帰納段階は，
　 $T(l + r) = T(l) + T(r) + l \times r \le cl^2 + cr^2 + 2clr = c(r^2 + l^2 + 2lr) = c(l + r)^2$
となるので $T(n) = O(n^2)$ ということが示せました．一般に分枝数が3以上でも同様の解析で示せます．
分枝数 $k$ で各子の部分木のサイズを $s_1, \ldots, s_k$ とすると上の漸化式から，
　 $T(s_1 + \cdots + s_k) = T(s_1) + \cdots + T(s_k) + s_1 s_2 + (s_1 + s_2) s_3 + \cdots + (s_1 + \cdots s_{k-1}) s_k$
となります．ここで，
　 $(s_1 + \cdots + s_k)^2 = ((s_1 + \cdots + s_{k-1}) + s_k)^2$
　　　　 $= (s_1 + \cdots + s_{k-1})^2 + s_k^2 + 2(s_1 + \cdots + s_{k-1})s_k$
　　　　 $= s_1^2 + \cdots + s_k^2 + 2s_1 s_2 + 2 (s_1 + s_2) s_3 + \cdots + 2(s_1 + \cdots + s_{k-1}) s_k$
となるので，
　 $T(s_1 + \cdots + s_k) \le c(s_1 + \cdots + s_k)^2$
となり， $T(n) = O(n^2)$ ということが示せました．

計算時間： $O(n^2)$

ソースコードを表示

#include <bits/stdc++.h>

using namespace std;

using ll = long long;
constexpr ll MOD = 1e9 + 7;

constexpr int MAX_N = 5050;
ll ch[MAX_N][2], dp[MAX_N][MAX_N][2];

class ProblemE {
public:
    ProblemE(int _n) : n(_n), m(n - 1), root(0), tree(n), comb(n + 1), sz(n, 1) {
        comb[2] = 1;
        for (int i = 4; i <= n; i += 2)
            comb[i] = ((i - 1) * comb[i - 2]) % MOD;
    }
    inline void add_edge(int v[2]) {
        tree[v[0] - 1].push_back(v[1] - 1);
        tree[v[1] - 1].push_back(v[0] - 1);
    }

    ll Solve() {
        Dfs(root, -1);

        ll res = 0;
        for (int i = 0; i <= n; ++i) {
            res = (res + (dp[root][i][0] * comb[i]) % MOD) % MOD;
            res = (res - (dp[root][i][1] * comb[i]) % MOD + MOD) % MOD;
        }

        return (res % MOD + MOD) % MOD;
    }

private:
    const int n, m, root;
    vector<vector<int>> tree;
    vector<ll> comb;
    vector<int> sz;

    void Dfs(const int cur, const int p) {
        dp[cur][1][0] = 1;

        for (const int dst : tree[cur]) {
            if (dst == p) continue;
            Dfs(dst, cur);
            for (int i = 0; i <= sz[cur] + sz[dst]; ++i)
                ch[i][0] = ch[i][1] = 0LL;

            for (int x = 1; x <= sz[cur]; ++x) {
                for (int y = 1; y <= sz[dst]; ++y) {
                    // combine cur with dst
                    (ch[x + y][0] += dp[cur][x][0] * dp[dst][y][0]) %= MOD;
                    (ch[x + y][0] += dp[cur][x][1] * dp[dst][y][1]) %= MOD;
                    (ch[x + y][1] += dp[cur][x][0] * dp[dst][y][1]) %= MOD;
                    (ch[x + y][1] += dp[cur][x][1] * dp[dst][y][0]) %= MOD;

                    // separate dst
                    (ch[x][1] += dp[cur][x][0] * dp[dst][y][0] % MOD * comb[y]) %= MOD;
                    (ch[x][1] += dp[cur][x][1] * dp[dst][y][1] % MOD * comb[y]) %= MOD;
                    (ch[x][0] += dp[cur][x][0] * dp[dst][y][1] % MOD * comb[y]) %= MOD;
                    (ch[x][0] += dp[cur][x][1] * dp[dst][y][0] % MOD * comb[y]) %= MOD;

                }
            }

            for (int i = 0; i <= sz[cur] + sz[dst]; ++i) {
                dp[cur][i][0] = ch[i][0];
                dp[cur][i][1] = ch[i][1];
            }

            sz[cur] += sz[dst];
        }
    }
};

int main() {
    cin.tie(0); ios::sync_with_stdio(false);

    int n, v[2];
    cin >> n;
    ProblemE p(n);
    for (int i = 0; i < n - 1; ++i) {
        cin >> v[0] >> v[1];
        p.add_edge(v);
    }
    cout << p.Solve() << endl;

    return 0;
}

まとめ

初めて「木 DP 二乗」を解いたので勉強になりました．上の実装で dp, ch をクラスのメンバー変数として std::vector で管理しようとすると TLE になります．グローバル変数で生配列が早いのは知っているのですが TLE になるほど遅いのは何ででしょうね．
std::vector を返す実装を多くの人がしていたので「木 DP 二乗」の実装に典型があるかもしれません．後で勉強します．

参考文献

[1] ABC #107 / ARC #101 Editorial（2018年8月28日アクセス）