|
| 1 | +// Rerooting |
| 2 | +// https://youtu.be/zG1L4vYuGrg?t=7092 |
| 3 | +// TODO: vertex info, edge info |
| 4 | +struct Rerooting { |
| 5 | + struct DP { |
| 6 | + DP() {} |
| 7 | + DP operator+(const DP& a) const { |
| 8 | + // edit here |
| 9 | + } |
| 10 | + DP addRoot() const { |
| 11 | + // edit here |
| 12 | + } |
| 13 | + }; |
| 14 | + |
| 15 | + int n; |
| 16 | + vector<vector<int>> to; |
| 17 | + vector<vector<DP>> dp; |
| 18 | + vector<DP> ans; |
| 19 | + Rerooting(int n=0):n(n),to(n),dp(n),ans(n) {} |
| 20 | + void addEdge(int a, int b) { |
| 21 | + to[a].push_back(b); |
| 22 | + to[b].push_back(a); |
| 23 | + } |
| 24 | + void init() { |
| 25 | + dfs(0); |
| 26 | + bfs(0); |
| 27 | + } |
| 28 | + |
| 29 | + DP dfs(int v, int p=-1) { |
| 30 | + DP dpSum; |
| 31 | + dp[v] = vector<DP>(to[v].size()); |
| 32 | + rep(i,to[v].size()) { |
| 33 | + int u = to[v][i]; |
| 34 | + if (u == p) continue; |
| 35 | + dp[v][i] = dfs(u,v); |
| 36 | + dpSum = dpSum + dp[v][i]; |
| 37 | + } |
| 38 | + return dpSum.addRoot(); |
| 39 | + } |
| 40 | + void bfs(int v, const DP& dpP=DP(), int p=-1) { |
| 41 | + int deg = to[v].size(); |
| 42 | + rep(i,deg) if (to[v][i] == p) dp[v][i] = dpP; |
| 43 | + |
| 44 | + vector<DP> dpSumL(deg+1); |
| 45 | + rep(i,deg) dpSumL[i+1] = dpSumL[i] + dp[v][i]; |
| 46 | + vector<DP> dpSumR(deg+1); |
| 47 | + for (int i = deg-1; i >= 0; --i) |
| 48 | + dpSumR[i] = dpSumR[i+1] + dp[v][i]; |
| 49 | + ans[v] = dpSumL[deg].addRoot(); |
| 50 | + |
| 51 | + rep(i,deg) { |
| 52 | + int u = to[v][i]; |
| 53 | + if (u == p) continue; |
| 54 | + DP d = dpSumL[i] + dpSumR[i+1]; |
| 55 | + bfs(u, d.addRoot(), v); |
| 56 | + } |
| 57 | + } |
| 58 | +}; |
0 commit comments