p3676 小清新数据结构题

分析

 

 树状数组做区间加时要用到差分的思想

 

 代码

#include<bits/stdc++.h>
using namespace std;
#define li long long
li Ans;
vector<int>v[200100];
inline int lb(int x){return x&(-x);}
int n,m,q,cnt,d1[200100],d2[200100],id[200100],w[200100],c[200100];
int dep[200100],fa[200100],son[200100],siz[200100],ac[200100],pl[200100];
inline void add(int x,int k){int wh=x*k;while(x<=n)d1[x]+=k,d2[x]+=wh,x+=lb(x);}
inline li qurey(int x){li res=0,wh=x+1;while(x)res+=1ll*wh*d1[x]-d2[x],x-=lb(x);return res;}
inline void dfs(int x,int f){
    int mx=0;
    siz[x]=1;
    dep[x]=dep[f]+1;
    fa[x]=f;
    for(int i=0;i<v[x].size();i++)
      if(v[x][i]!=f){
          dfs(v[x][i],x);
          w[x]+=w[v[x][i]];
          siz[x]+=siz[v[x][i]];
          if(siz[v[x][i]]>mx){
            mx=siz[v[x][i]];
            son[x]=v[x][i];
        }
      }
    Ans+=1ll*w[x]*w[x];
}
inline void dfs2(int x,int acc){
    ac[x]=acc;
    id[x]=++cnt;
    pl[cnt]=x;
    add(cnt,w[x]-w[pl[cnt-1]]);
    if(son[x])dfs2(son[x],acc);
    for(int i=0;i<v[x].size();i++)
      if(v[x][i]!=fa[x]&&v[x][i]!=son[x])dfs2(v[x][i],v[x][i]);
}
inline void update(int x,int k){
    int sum=0;
    li res=0;
    while(x){
      sum+=id[x]-id[ac[x]]+1;
      res+=qurey(id[x])-qurey(id[ac[x]]-1);
      add(id[ac[x]],k),add(id[x]+1,-k);
      x=fa[ac[x]];
    }
    Ans+=1ll*k*(sum*k+(res<<1));
}
inline li que(int x){
    int sum=0;
    li res=0,k=qurey(1);
    while(x){
      sum+=id[x]-id[ac[x]]+1;
      res+=qurey(id[x])-qurey(id[ac[x]]-1);
      x=fa[ac[x]];
    }
    return Ans+k*((sum+1)*k-(res<<1));
}
int main(){
    int i,j,k;
    scanf("%d%d",&n,&q);
    for(i=1;i<n;i++){
      int x,y;
      scanf("%d%d",&x,&y);
      v[x].push_back(y);
      v[y].push_back(x);
    }
    for(i=1;i<=n;i++)scanf("%d",&c[i]),w[i]=c[i];
    dfs(1,0);
    dfs2(1,1);
    for(i=1;i<=q;i++){
      int x,y;
      scanf("%d%d",&k,&x);
      if(k==1)scanf("%d",&y),y-=c[x],c[x]+=y,update(x,y);
        else printf("%lld\n",que(x));
    } 
    return 0;
}