給一顆樹,每個點上有權值。有兩種操作,修改某個點的權值,或者查詢距離某個點不超過d的點的權值和。
數據範圍:樹的規模和操作次數均爲10^5。
樹的分治,對於每個分治中心以及他們的孩子,維護樹狀數組,距離他們d的點的權值和。
#pragma comment(linker, "/STACK:1024000000,1024000000")
#include <cstdio>
#include <algorithm>
#include <vector>
#include <iostream>
using namespace std;
struct BIT {
vector<int> a;
void clear(int n) {
a.clear();
a.resize(n+1,0);
}
int lb(int x) {
return x&-x;
}
int get(int i) {
if (i>=a.size()) i=a.size()-1;
int ans=0;
for (;i>0;i-=lb(i)) ans+=a[i];
return ans;
}
void set(int i,int x) {
for (;i<a.size();i+=lb(i)) a[i]+=x;
}
void add(BIT &b) {
int last=0;
for (int i=1;i<b.a.size();i++) {
int tmp=b.get(i);
set(i,tmp-last);
last=tmp;
}
}
};
struct Node {
int belong[20],from[20],dis[20];
BIT num[20];
int v,fe,n;
bool visited;
void clear() {
fe=-1;
visited=false;
}
};
struct Edge {
int t,ne;
};
Node a[100001];
Edge b[200000];
int bp;
void putedge(int x,int y) {
b[bp].t=y;
b[bp].ne=a[x].fe;
a[x].fe=bp++;
}
int center,centerV;
void getCenter(int i,int n) {
a[i].visited=true;
int maxn=0;
a[i].n=1;
for (int j=a[i].fe;j!=-1;j=b[j].ne) {
if (!a[b[j].t].visited) {
getCenter(b[j].t,n);
a[i].n+=a[b[j].t].n;
maxn=max(maxn,a[b[j].t].n);
}
}
maxn=max(maxn,n-a[i].n);
if (maxn<centerV) {
centerV=maxn;
center=i;
}
a[i].visited=false;
}
void calNum(int i,int kk,int from,int belong,int dis) {
a[i].visited=true;
a[i].dis[kk]=dis;
a[i].belong[kk]=belong;
a[i].from[kk]=from;
a[from].num[kk].set(dis,a[i].v);
for (int j=a[i].fe;j!=-1;j=b[j].ne) {
if (!a[b[j].t].visited) {
calNum(b[j].t,kk,from,belong,dis+1);
}
}
a[i].visited=false;
}
void calAns(int i,int n,int kk) {
centerV=10000000;
getCenter(i,n);
i=center;
a[i].num[kk].clear(centerV+1);
getCenter(i,n);
a[i].visited=true;
a[i].belong[kk]=i;
for (int j=a[i].fe;j!=-1;j=b[j].ne) {
if (!a[b[j].t].visited) {
a[b[j].t].num[kk].clear(a[b[j].t].n+1);
calNum(b[j].t,kk,b[j].t,i,1);
a[i].num[kk].add(a[b[j].t].num[kk]);
}
}
for (int j=a[i].fe;j!=-1;j=b[j].ne)
if (!a[b[j].t].visited)
calAns(b[j].t,a[b[j].t].n,kk+1);
}
void change(int x,int y) {
int k=0;
while (a[x].belong[k]!=x) {
a[a[x].belong[k]].num[k].set(a[x].dis[k],y-a[x].v);
a[a[x].from[k]].num[k].set(a[x].dis[k],y-a[x].v);
k++;
}
a[x].v=y;
}
int get(int x,int y) {
int k=0,ans=0;
while (a[x].belong[k]!=x) {
int dis=a[x].dis[k];
int i=a[x].belong[k];
int j=a[x].from[k];
if (y>dis) {
ans+=a[i].num[k].get(y-dis);
ans-=a[j].num[k].get(y-dis);
}
if (y>=dis) ans+=a[i].v;
k++;
}
ans+=a[x].num[k].get(y);
ans+=a[x].v;
return ans;
}
void print(int n) {
int i,j,k;
printf("PRINT\n");
for (i=1;i<=n;i++) {
printf("Node %d:\n",i);
printf(" belong:");
for (j=0;j<3;j++) printf(" %d",a[i].belong[j]);
printf("\n from:");
for (j=0;j<3;j++) printf(" %d",a[i].from[j]);
printf("\n dis:");
for (j=0;j<3;j++) printf(" %d",a[i].dis[j]);
printf("\n num:");
for (j=0;j<3;j++) {
printf("\n %d:",j);
for (k=1;k<a[i].num[j].a.size();k++)
printf(" %d",a[i].num[j].get(k)-a[i].num[j].get(k-1));
}
printf("\n");
}
}
int main() {
int n,q,i,cas=1;
while (scanf("%d%d",&n,&q)!=EOF) {
//cerr << "Case: " << cas++ << endl;
for (i=1;i<=n;i++) {
a[i].clear();
scanf("%d",&a[i].v);
}
bp=0;
for (i=1;i<n;i++) {
int x,y;
scanf("%d%d",&x,&y);
putedge(x,y);
putedge(y,x);
}
calAns(1,n,0);
for (i=0;i<q;i++) {
char c;
int x,y;
scanf(" %c%d%d",&c,&x,&y);
//cerr << cas << ' ' << i << endl;
//cerr << c << ' ' << x << ' ' << y << endl;
if (c=='!') change(x,y);
else printf("%d\n",get(x,y));
//print(n);
}
}
return 0;
}