Educational DP Contest V問題：Subtree

問題. Subtree

$n$ 頂点の木 $T$ がある．各頂点 $v \in V(T)$ に対して， $v$ を含む部分木の数を答えよ．

制約： $1 \le n \le 10^5$

解法. 全方位木DP

解法が分からなかったのでkmjpさんとsnukeさんの解説を参考にしました．

$f(p, v)$ を， $v$ の親を $p$ とした根付き木で $v$ を含み $v$ を根とする部分木の数　として動的計画法を行います．頂点 $v$ の隣接頂点全体を $N(v)$ とすると遷移は，

　　 $f(p, v) = \prod_{c \in N(v) \setminus \{p\}} \left(f(v, c) + 1 \right)$ 　　（☆）

となります． $f(v, c) + 1$ となっているのは，頂点 $c$ を含まない空の部分木の数も考慮しているからです．もし，各辺 $\{u, v\} \in E(T)$ に対して， $f(u, v)$ と $f(v, u)$ が計算できているとすると各頂点 $v$ に対する答えは

　　 $\prod_{c \in N(v)} \left(f(v, c) + 1 \right)$

となります．この出力は $O(n)$ 時間でできるので，ここからは $f(u, v)$ と $f(v, u)$ を $O(n)$ 時間で求めることを目標に説明していきます．

適当な頂点 $r \in V(T)$ を選び， $r$ を根とする順序木 $T_v$ を考えます．順序木とは根付き木で各頂点 $v$ の子全体 $C(v)$ に全順序 $\le_{v}$ を与えたものです． $r$ から深さ優先探索で上の漸化式（☆）を計算することによって親 $p$ からその子 $c$ への値 $f(p, c)$ を $O(n)$ 時間で求めることができます．しかし逆向きの $f(c, p)$ は計算されていないためにすべての頂点の答えを計算することはできません．下の図では赤色と青色の矢印が計算したいもので，実線の赤色矢印は計算されていますが青色の点線矢印はまだ計算されていません．
　　　　 f:id:tatanaideyo:20190129161340p:plain:w450

各頂点を根として同様の計算をすることによって青色の矢印の値を計算できますが $O(n^2)$ 時間となり間に合いません．ここで， $f(v_1, r)$ を計算することを考えます．漸化式（☆）から，　
　　 $f(v_1, r) = (f(r, v_2) + 1) \times (f(r, v_3) + 1) \times (f(r, v_4) + 1)$
となります．右辺の項はすでに計算されているために $|C(r)| - 1$ 回の積で求めることができます．一般的に頂点 $p$ へ向かう青色の矢印は子 $c \in C(p)$ に対して，
　　 $f(c, p) = \prod_{v \in N(p) \setminus \{c\}} \left(f(p, v) + 1 \right)$
と計算できます．各 $f(p, v)$ は赤色の矢印ですでに計算されているのですが， $p$ へ向かうすべての青色の矢印を愚直に計算すると $(|C(v)| - 1)^2$ 回の積が必要となります．したがって全体で $O(n^2)$ 時間かかってしまうのでこの部分を高速にする必要があります．

頂点 $p$ の子を $c_1, c_2, \ldots, c_k$ とします．ここで任意の $1 \le i \le k$ に対して，
　　 $g_l(i) = \left( f(p, c_1) + 1 \right) \times \left( f(p, c_2) + 1 \right) \times \cdots \times \left( f(p, c_i) + 1 \right)$ ,
　　 $g_r(i) = (f(p, c_k) + 1) \times (f(p, c_{k - 1}) + 1) \times \cdots \times (f(p, c_i) + 1)$
とします．ただし， $g_l(0)$ と $g_r(k + 1)$ の値は 1 とします．これを前処理で $O(|C(p)|)$ 時間で求めることによって，各子 $c_i$ に対して，
　　 $f(c_i, p) = g_l(i - 1) \times g_r(i + 1)$
と定数時間で求めることができます．したがって， $p$ へ向かうすべての青色の矢印を $O(|C(p)|)$ 時間で求めることができるので，全体で $O(n)$ 時間で答えを計算することができます．下の図は $f(c_2, p)$ の計算を表しています．
　　　　
　　　　 f:id:tatanaideyo:20190129172705p:plain:w450

計算時間： $O(n)$

ソースコードを表示

#include <bits/stdc++.h>

using namespace std;

using ll = long long;

int N, MOD;
vector<vector<int>> adj;
vector<vector<ll>> dp;

ll Dfs(const int p, const int v) {
    ll res = 1;
    for (size_t i = 0; i < adj[v].size(); ++i) {
        if (adj[v][i] != p) {
            dp[v][i] = Dfs(v, adj[v][i]);
            (res *= (dp[v][i] + 1)) %= MOD;
        }
    }
    return res;
}

void Rev(const int p, const int v, const ll add) {
    for (size_t i = 0; i < adj[v].size(); ++i)
        if (adj[v][i] == p) {
            dp[v][i] = add;
            break;
        }

    const int deg = adj[v].size();
    vector<ll> l(deg), r(deg);
    for (int i = 0; i < deg; ++i) r[i] = ((l[i] = (dp[v][i] + 1)) %= MOD);
    for (int i = 1; i < deg; ++i) (l[i] *= l[i - 1]) %= MOD;
    for (int i = deg - 2; 0 <= i; --i) (r[i] *= r[i + 1]) %= MOD;

    for (int i = 0; i < deg; ++i) {
        if (adj[v][i] == p) continue;

        ll add_c = 1;
        if (i != 0) (add_c *= l[i - 1]) %= MOD;
        if (i != deg - 1) (add_c *= r[i + 1]) %= MOD;
        Rev(v, adj[v][i], add_c);
    }
}

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

    cin >> N >> MOD;

    adj.resize(N);
    for (int i = 0, x, y; i < N - 1; ++i) {
        cin >> x >> y;
        adj[x - 1].push_back(y - 1);
        adj[y - 1].push_back(x - 1);
    }

    dp.resize(N);
    for (int v = 0; v < N; ++v) dp[v].resize(adj[v].size());
    Dfs(-1, 0);
    Rev(-1, 0, 0);

    for (int v = 0; v < N; ++v) {
        ll ans = 1;
        for (ll add : dp[v]) (ans *= (add + 1)) %= MOD;
        cout << ans << '\n';
    }

    return 0;
}

忘れても大丈夫

Educational DP Contest V問題：Subtree

問題. Subtree

解法. 全方位木DP

参考文献