Submission #11405065


Source Code Expand

Copy
/* rerooting.cpp
    全方位木DP

    // 根rから最も遠い点までの距離を求める
    struct DP {  // 単位元はしっかり定義する(末端でもfrされるので注意)
        long long dp;
        DP(long long dp_ = 0) : dp(dp_) {}
    };
    function<DP(DP, DP, long long)> fd = [](DP dp_cum, DP d, long long cost) -> DP {  // d:辺eに対応する部分木のdpの値  cost:eのコスト
        return DP(max(dp_cum.dp, d.dp + cost));                                       //最大の深さ(根から最も遠い点までの距離)
    };
    function<DP(DP)> fr = [](DP d) -> DP {  // まとめたDPを用いて、新たな部分木のDPを計算する
        return d;
    };

    verified:
    - AOJ GRL_5_A: Diameter of a Tree
        http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=GRL_5_A
    - AOJ 1595: Traffic Tree
        http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=1595
*/

#include <bits/stdc++.h>
using namespace std;

template <int mod>
struct ModInt {
    int val;
    ModInt() : val(0) {}
    ModInt(long long x) : val(x >= 0 ? x % mod : (mod - (-x) % mod) % mod) {}
    int getmod() { return mod; }
    ModInt &operator+=(const ModInt &p) {
        if ((val += p.val) >= mod) {
            val -= mod;
        }
        return *this;
    }
    ModInt &operator-=(const ModInt &p) {
        if ((val += mod - p.val) >= mod) {
            val -= mod;
        }
        return *this;
    }
    ModInt &operator*=(const ModInt &p) {
        val = (int)(1LL * val * p.val % mod);
        return *this;
    }
    ModInt &operator/=(const ModInt &p) {
        *this *= p.inverse();
        return *this;
    }
    ModInt operator-() const { return ModInt(-val); }
    ModInt operator+(const ModInt &p) const { return ModInt(*this) += p; }
    ModInt operator-(const ModInt &p) const { return ModInt(*this) -= p; }
    ModInt operator*(const ModInt &p) const { return ModInt(*this) *= p; }
    ModInt operator/(const ModInt &p) const { return ModInt(*this) /= p; }
    bool operator==(const ModInt &p) const { return val == p.val; }
    bool operator!=(const ModInt &p) const { return val != p.val; }
    ModInt inverse() const {
        int a = val, b = mod, u = 1, v = 0, t;
        while (b > 0) {
            t = a / b;
            swap(a -= t * b, b);
            swap(u -= t * v, v);
        }
        return ModInt(u);
    }
    ModInt pow(long long n) const {
        ModInt ret(1), mul(val);
        while (n > 0) {
            if (n & 1) ret *= mul;
            mul *= mul;
            n >>= 1;
        }
        return ret;
    }
    friend ostream &operator<<(ostream &os, const ModInt &p) { return os << p.val; }
    friend istream &operator>>(istream &is, ModInt &a) {
        long long t;
        is >> t;
        a = ModInt<mod>(t);
        return (is);
    }
    static int get_mod() { return mod; }
};
const int MOD = 1000000007;  // if inv is needed, this shold be prime.
using modint = ModInt<MOD>;

/* Comb:modintで二項係数を計算する構造体
    前処理:O(n)
    二項係数の計算:O(1)
    制約:
        n<=10^7
        k<=10^7
        p(mod)は素数
*/
template <class T>
struct Comb {
    vector<T> fact_, fact_inv_, inv_;
    // Comb() {}
    Comb(int SIZE) : fact_(SIZE, 1), fact_inv_(SIZE, 1), inv_(SIZE, 1) { init(SIZE); }
    void init(int SIZE) {
        fact_.assign(SIZE, 1), fact_inv_.assign(SIZE, 1), inv_.assign(SIZE, 1);
        int mod = fact_[0].getmod();
        for (int i = 2; i < SIZE; i++) {
            fact_[i] = fact_[i - 1] * i;
            inv_[i] = -inv_[mod % i] * (mod / i);
            fact_inv_[i] = fact_inv_[i - 1] * inv_[i];
        }
    }
    T nCk(int n, int k) {
        assert(!(n < k));
        assert(!(n < 0 || k < 0));
        return fact_[n] * fact_inv_[k] * fact_inv_[n - k];
    }
    T nHk(int n, int k) {
        assert(!(n < 0 || k < 0));
        return nCk(n + k - 1, k);
    }
    T fact(int n) {
        assert(!(n < 0));
        return fact_[n];
    }
    T fact_inv(int n) {
        assert(!(n < 0));
        return fact_inv_[n];
    }
    T inv(int n) {
        assert(!(n < 0));
        return inv_[n];
    }
};

Comb<modint> comb(10000);

/* Rerooting: 全方位木 DP
    問題ごとに以下を書き換える
    - 型DPと単位元
    - 型DPに対する二項演算 fd
    - まとめたDPを用いて新たな部分木のDPを計算する fr
    計算量: O(N)
*/
struct Rerooting {
    /* start 問題ごとに書き換え */
    struct DP {  // DP の型
        modint dp;
        int s;
        DP(modint dp_, int s_) : dp(dp_), s(s_) {}
    };
    const DP unit_dp = DP(modint(1), 0);                                              // 単位元はしっかり定義する(末端でもfrされるので注意)
    function<DP(DP, DP, long long)> fd = [](DP dp_cum, DP d, long long cost) -> DP {  // d:辺eに対応する部分木のdpの値  cost:eのコスト
        int n = dp_cum.s + d.s;
        return DP(dp_cum.dp * d.dp * comb.nCk(n, d.s), n);
    };
    function<DP(DP)> fr = [](DP d) -> DP {  // まとめたDPを用いて、新たな部分木のDPを計算する
        return DP(d.dp, d.s + 1);
    };
    /* end 問題ごとに書き換え */

    // グラフの定義
    struct Edge {
        int from;
        int to;
        long long cost;
    };
    using Graph = vector<vector<Edge>>;

    vector<vector<DP>> dp;  // dp[v][i]: vから出るi番目の有向辺に対応する部分木のDP
    vector<DP> ans;         // ans[v]: 頂点vを根とする木の答え
    Graph G;

    Rerooting(int N) : G(N) {
        dp.resize(N);
        ans.assign(N, unit_dp);
    }

    void add_edge(int a, int b, long long c = 1) {
        G[a].push_back({a, b, c});
    }
    void build() {
        dfs(0);           // 普通に木DP
        bfs(0, unit_dp);  // 残りの部分木に対応するDPを計算
    }

    DP dfs(int v, int p = -1) {  // 頂点v, 親p
        DP dp_cum = unit_dp;
        int deg = G[v].size();
        dp[v] = vector<DP>(deg, unit_dp);
        for (int i = 0; i < deg; i++) {
            int u = G[v][i].to;
            if (u == p) continue;
            dp[v][i] = dfs(u, v);
            dp_cum = fd(dp_cum, dp[v][i], G[v][i].cost);
        }
        return fr(dp_cum);
    }
    void bfs(int v, const DP &dp_p, int p = -1) {
        int deg = G[v].size();
        for (int i = 0; i < deg; i++) {  // 前のbfsで計算した有向辺に対応する部分木のDPを保存
            if (G[v][i].to == p) dp[v][i] = dp_p;
        }

        vector<DP> dp_l(deg + 1, unit_dp), dp_r(deg + 1, unit_dp);  // 累積的なDP
        for (int i = 0; i < deg; i++) {
            dp_l[i + 1] = fd(dp_l[i], dp[v][i], G[v][i].cost);
        }
        for (int i = deg - 1; i >= 0; i--) {
            dp_r[i] = fd(dp_r[i + 1], dp[v][i], G[v][i].cost);
        }
        ans[v] = fr(dp_l[deg]);
        for (int i = 0; i < deg; i++) {
            int u = G[v][i].to;
            if (u == p) continue;
            bfs(u, fr(fd(dp_l[i], dp_r[i + 1], 0)), v);  // 累積同士のfdなので、edgeは適当に
        }
    }
};

int main() {
    int N;
    cin >> N;
    Rerooting reroot(N);
    for (int i = 0; i < N - 1; i++) {
        int u, v;
        cin >> u >> v;
        u--, v--;
        reroot.add_edge(u, v);
        reroot.add_edge(v, u);
    }
    reroot.build();

    modint ans = (0);
    for (int i = 0; i < N; i++) {
        ans += reroot.ans[i].dp;
    }
    cout << ans / 2 << endl;
}

Submission Info

Submission Time
Task N - 木
User kami634
Language C++14 (GCC 5.4.1)
Score 5
Code Size 7745 Byte
Status AC
Exec Time 3 ms
Memory 896 KB

Judge Result

Set Name All
Score / Max Score 5 / 5
Status
AC × 9
Set Name Test Cases
All 00, 01, 02, 03, 04, 05, 06, 90, 91
Case Name Status Exec Time Memory
00 AC 3 ms 896 KB
01 AC 2 ms 512 KB
02 AC 2 ms 512 KB
03 AC 2 ms 512 KB
04 AC 2 ms 512 KB
05 AC 2 ms 512 KB
06 AC 2 ms 512 KB
90 AC 1 ms 384 KB
91 AC 1 ms 384 KB