Submission #11405065
Source Code Expand
/* 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 |
|
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 |