-
Notifications
You must be signed in to change notification settings - Fork 0
/
rerooting.hpp
64 lines (62 loc) · 1.66 KB
/
rerooting.hpp
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
#pragma once
template<typename M,typename N,M (*f)(M,N),M (*g)(M,int),M (*h)(M,M),M (*e)()>struct Rerooting{
using P=pair<int,N>;
vector<vector<P>> G;
vector<M> buf,ret;
Rerooting(int n):G(n){}
void add_edge(int u,int v,N w){
G[u].push_back({v,w});
G[v].push_back({u,w});
}
void dfs1(int v,int p){
for(auto& [to,w]:G[v])if(to!=p){
dfs1(to,v);
buf[v]=h(buf[v],f(buf[to],w));
}
buf[v]=g(buf[v],v);
}
void dfs2(int v,int p,M pc){
buf[v]=e();
vector<M> cs;
for(auto& [to,w]:G[v])if(to!=p){
cs.push_back(f(buf[to],w));
buf[v]=h(buf[v],f(buf[to],w));
}
buf[v]=g(h(buf[v],pc),v);
ret[v]=buf[v];
cs.push_back(pc);
int m=cs.size();
vector<M> L(m),R(m);
rep(i,0,m){
L[i]=cs[i];
if(i)L[i]=h(L[i],L[i-1]);
}
for(int i=m-1;i>=0;i--){
R[i]=cs[i];
if(i!=m-1)R[i]=h(R[i],R[i+1]);
}
int id=0;
for(auto& [to,w]:G[v])if(to!=p){
M nxt=e();
if(id!=0)nxt=h(nxt,L[id-1]);
if(id!=m-1)nxt=h(nxt,R[id+1]);
M buf_v=buf[v],buf_to=buf[to];
buf[v]=g(nxt,v);
M add=f(buf[v],w);
dfs2(to,v,add);
buf[v]=buf_v,buf[to]=buf_to;
id++;
}
}
vector<M> run(){
buf.assign(G.size(),e());
ret.assign(G.size(),e());
dfs1(0,-1);
dfs2(0,-1,e());
return ret;
}
};
/**
* @brief Rerooting
* @docs docs/rerooting.md
*/