[树链剖分][树状数组] Luogu P3676 小清新数据结构题
题目描述
在很久很久以前,有一棵n个点的树,每个点有一个点权。
现在有q次操作,每次操作是修改一个点的点权或指定一个点,询问以这个点为根时每棵子树点权和的平方和。
(题目不是很好懂,没看太懂的可以看看样例解释)
题解
- 大爷博客(戳我)
代码
1 #include <cstdio> 2 #include <cstring> 3 #include <iostream> 4 #define ll long long 5 #define sqr(x) (x)*(x) 6 using namespace std; 7 const ll N=2e5+10; 8 struct edge{ll from,to;}e[N<<1]; 9 ll n,m,cnt,ans,tot,w[N],c[N],fa[N],sz[N],id[N],pos[N],deep[N],son[N],top[N],head[N],sz1[N],sz2[N]; 10 void add(ll x,ll y) { for (int i=x;i<=n;i+=i&(-i)) sz1[i]+=y,sz2[i]+=x*y; } 11 ll query(ll x) { ll r=0; for (ll i=x;i;i-=i&(-i)) r+=(x+1)*sz1[i]-sz2[i]; return r; } 12 void insert(ll x,ll y) 13 { 14 e[++cnt].to=y,e[cnt].from=head[x],head[x]=cnt; 15 e[++cnt].to=x,e[cnt].from=head[y],head[y]=cnt; 16 } 17 void dfs1(ll x) 18 { 19 deep[x]=deep[fa[x]]+1,sz[x]=1; 20 for (ll i=head[x];i;i=e[i].from) if (e[i].to!=fa[x]) 21 { 22 fa[e[i].to]=x,dfs1(e[i].to),sz[x]+=sz[e[i].to],w[x]+=w[e[i].to]; 23 if (sz[e[i].to]>sz[son[x]]) son[x]=e[i].to; 24 } 25 ans+=sqr(w[x]); 26 } 27 void dfs2(ll x,ll pre) 28 { 29 top[x]=pre;pos[id[x]=++tot]=x,add(id[x],w[x]-w[pos[tot-1]]); 30 if (son[x]) dfs2(son[x],pre); 31 for (ll i=head[x];i;i=e[i].from) if (fa[x]!=e[i].to&&e[i].to!=son[x]) dfs2(e[i].to,e[i].to); 32 } 33 void modify(ll x,ll y) 34 { 35 ll l=0,r=0; 36 for (;x;x=fa[top[x]]) l+=id[x]-id[top[x]]+1,r+=query(id[x])-query(id[top[x]]-1),add(id[top[x]],y),add(id[x]+1,-y); 37 ans+=y*(y*l+(r<<1)); 38 } 39 ll calc(ll u) 40 { 41 ll k=0,x=query(1),s=0; 42 for (;u;u=fa[top[u]]) k+=id[u]-id[top[u]]+1,s+=query(id[u])-query(id[top[u]]-1); 43 return ans+x*((k+1)*x-(s<<1)); 44 } 45 int main() 46 { 47 scanf("%lld%lld",&n,&m); 48 for (ll i=2,x,y;i<=n;i++) scanf("%lld%lld",&x,&y),insert(x,y); 49 for (ll i=1;i<=n;i++) scanf("%lld",&c[i]),w[i]=c[i]; 50 dfs1(1),dfs2(1,1); 51 for (ll op,x,y;m;m--) 52 { 53 scanf("%lld%lld",&op,&x); 54 if (op==1) scanf("%lld",&y),y-=c[x],c[x]+=y,modify(x,y); else printf("%lld\n",calc(x)); 55 } 56 }