-
Notifications
You must be signed in to change notification settings - Fork 0
/
rerootingDP.cc
68 lines (61 loc) · 1.73 KB
/
rerootingDP.cc
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
template <class S, class V, class E,
S (*opS)(S, S),
S (*opV)(S, V),
S (*opE)(S, E),
S (*e)()>
struct RerootingDP {
int n;
vector<vector<int>> adj;
vector<V> vertex_info;
vector<vector<E>> edge_info;
vector<vector<S>> dp;
vector<S> ans;
RerootingDP(int n_) : n(n_), adj(n_), vertex_info(n_), edge_info(n_), dp(n_), ans(n_, e()) {}
void add_edge(int from, int to, const E& info=E()) {
adj[from].push_back(to);
edge_info[from].emplace_back(info);
};
void solve() {
/* O(N) */
dfs(0);
reroot(0);
}
private:
S dfs(int v, int p=-1) {
int deg = ssize(adj[v]);
dp[v] = vector<S>(deg, e());
S x = e();
for(int i=0; i<deg; i++) {
int u = adj[v][i];
if(u==p) continue;
dp[v][i] = opE(dfs(u,v), edge_info[v][i]);
x = opS(x, dp[v][i]);
}
return opV(x, vertex_info[v]);
}
void reroot(int v, int p=-1, const S& pval=e()) {
int deg = ssize(adj[v]);
for(int i=0; i<deg; i++) if(adj[v][i] == p) {
dp[v][i] = pval;
break;
}
vector<S> sumL(deg+1, e()), sumR(deg+1, e());
for(int i=0; i<deg; i++) sumL[i+1] = opS(sumL[i], dp[v][i]);
for(int i=deg-1; i>=0; i--) sumR[i] = opS(sumR[i+1], dp[v][i]);
ans[v] = sumL[deg];
for(int i=0; i<deg; i++) {
int u = adj[v][i];
if (u == p) continue;
reroot(u, v, opE(opV(opS(sumL[i], sumR[i+1]), vertex_info[v]), edge_info[v][i]));
}
}
};
#if 0
struct V { V() {}; };
struct E { E() {}; };
struct S { S() {}; };
S e() { return ; }
S opS(S a, S b) { /* merge child nodes */ return ; }
S opV(S s, V info) { /* apply V when send value to the parent */ return s; }
S opE(S s, E info) { /* apply E when recieving value from a child */ return s; }
#endif