hdu_4918_Query on the subtree(树的分治+树状数组)

  1 #include<bits/stdc++.h>
  2 #define F(i,a,b) for(int i=a;i<=b;i++)
  3 #define pb push_back
  4 using namespace std;
  5 const int N=1e5+7;
  6 
  7 int n,q,g[N],nxt[N*2],v[N*2],ed,w[N],vis[N],id[N];
  8 int pool[40*N],C_ed,pool_ed,sz[N],mi,mx[N],ROOT;
  9 char op[2];
 10 
 11 struct node
 12 {
 13     int rt,subrt,dis;
 14     node(){}
 15     node(int _rt,int _subrt,int _dis):rt(_rt),subrt(_subrt),dis(_dis){}
 16 }tmp;
 17 vector<node>vt[N];
 18 
 19 void adg(int x,int y){v[++ed]=y,nxt[ed]=g[x],g[x]=ed;}
 20 void init(){ed=C_ed=pool_ed=0;F(i,1,n)vt[i].clear(),vis[i]=g[i]=0;}
 21 inline void up(int &a,int b){if(a<b)a=b;}
 22 
 23 struct BIT
 24 {
 25     int *C,n;
 26     void init(int tot){n=tot,C=pool+pool_ed,pool_ed+=tot+1;F(i,0,n)C[i]=0;}
 27     inline void add(int x,int c){while(x<=n)C[x]+=c,x+=x&-x;}
 28     inline int ask(int x,int an=0)
 29     {
 30         if(x>n)x=n;
 31         while(x>0)an+=C[x],x-=x&-x;
 32         return an;
 33     }
 34 }tr[N*2];
 35 
 36 void get_rt(int u,int fa,int num)
 37 {
 38     sz[u]=1,mx[u]=0;
 39     for(int i=g[u];i;i=nxt[i])
 40         if(!vis[v[i]]&&v[i]!=fa)
 41         {
 42             get_rt(v[i],u,num);
 43             sz[u]+=sz[v[i]],up(mx[u],sz[v[i]]);
 44         }
 45     up(mx[u],num-sz[u]);
 46     if(mx[u]<mi)ROOT=u,mi=mx[u];
 47 }
 48 
 49 void del(int u,int fa,int rt,int subrt,int dis=1)//将子树的每个点放进对应的重心
 50 {
 51     vt[u].pb(node(rt,subrt,dis));
 52     tr[rt].add(dis+1,w[u]);
 53     tr[subrt].add(dis+1,w[u]);
 54     for(int i=g[u];i;i=nxt[i])
 55         if(v[i]!=fa&&!vis[v[i]])
 56             del(v[i],u,rt,subrt,dis+1);
 57 }
 58 
 59 void init_tree(int u=1,int num=n)
 60 {
 61     mi=N,get_rt(u,u,num);
 62     int rt=ROOT,rt_id=++C_ed;
 63     tr[C_ed].init(sz[u]+2);
 64     vis[rt]=1,vt[rt].pb(node(C_ed,0,0));
 65     tr[C_ed].add(1,w[rt]);
 66     get_rt(rt,rt,num);//从新计算sz的大小
 67     for(int i=g[rt];i;i=nxt[i])
 68         if(!vis[v[i]])
 69         {
 70             tr[++C_ed].init(sz[v[i]]+2);
 71             del(v[i],v[i],rt_id,C_ed);
 72         }
 73     for(int i=g[rt];i;i=nxt[i])
 74         if(!vis[v[i]])
 75             init_tree(v[i],sz[v[i]]);
 76 }
 77 
 78 int main()
 79 {
 80     while(~scanf("%d%d",&n,&q))
 81     {
 82         init();
 83         F(i,1,n)scanf("%d",w+i);
 84         F(i,1,n-1)
 85         {
 86             int x,y;
 87             scanf("%d%d",&x,&y);
 88             adg(x,y),adg(y,x);
 89         }
 90         init_tree();
 91         while(q--)
 92         {
 93             int u,v;
 94             scanf("%s%d%d",op,&u,&v);
 95             if(op[0]==‘!‘)
 96             {
 97                 int size=vt[u].size(),d=v-w[u];
 98                 F(i,0,size-1)
 99                 {
100                     tmp=vt[u][i];
101                     tr[tmp.rt].add(tmp.dis+1,d);//dis+1去掉距离为0在BIT上超时的BUG
102                     if(tmp.subrt)tr[tmp.subrt].add(tmp.dis+1,d);
103                 }
104                 w[u]+=d;
105             }else
106             {
107                 int d=v,ans=0,size=vt[u].size();
108                 F(i,0,size-1)
109                 {
110                     tmp=vt[u][i];
111                     ans+=tr[tmp.rt].ask(d-tmp.dis+1);
112                     if(tmp.subrt)ans-=tr[tmp.subrt].ask(d-tmp.dis+1);
113                 }
114                 printf("%d\n",ans);
115             }
116         }
117     }
118     return 0;
119 }