Query on the subtree
Time Limit: 16000/8000 MS (Java/Others) Memory Limit: 131072/131072 K (Java/Others)
Total Submission(s): 536 Accepted Submission(s): 180
i
.
There are q operations. Each operations are of the following 2 types:
Change the weight of vertex v into x (denoted as “! v x”),
Ask the total weight of vertices whose distance are no more than d away from vertex v (denoted as “? v d”).
Note that the distance between vertex u and v is the number of edges on the shortest path between them.
The first line contains n,q (1≤n,q≤10
5
). The second line contains n integers w
1
,w
2
,…,w
n
(0≤w
i
≤10
4
). Each of the following (n – 1) lines contain 2 integers a
i
,b
i
denoting an edge between vertices a
i
and b
i
(1≤a
i
,b
i
≤n). Each of the following q lines contain the operations (1≤v≤n,0≤x≤10
4
,0≤d≤n).
For each queries, a single number denotes the total weight.
4 3 1 1 1 1 1 2 2 3 3 4 ? 2 1 ! 1 0 ? 2 1 3 3 1 2 3 1 2 1 3 ? 1 0 ? 1 1 ? 1 2
3 2 1 6 6
题意:给定一棵树,有两种操作,1.修改某个节点的值;2.询问与某个点距离不大于d的节点的权值和是多少
思路:树分治+树状数组,树分治的过程中记录下每个节点对应的重心,他是在这个中心的哪个子树中,他距离中心的距离是多少,树状数组记录这个节点能被哪些节点更新到
树分治的关键是每个节点最多被logn个节点更新到
#include<iostream>
#include<cstdio>
#include<string>
#include<cstring>
#include<vector>
#include<cmath>
#include<queue>
#include<stack>
#include<map>
#include<set>
#include<algorithm>
using namespace std;
const int maxn=100010;
const int INF=1000000000;
int N,Q;
int cnt,id[maxn];
int w[maxn];
int pool[maxn*40];//树状数组缓冲池
int *tail;
int tot,head[maxn];
int root,Max;
int size[maxn],maxv[maxn];
struct BIT
{
int *tree;
int n;
void init(int tot)
{
n=tot;
tree=tail;
tail=tail+tot;
memset(tree,0,sizeof(int)*n);
}
void add(int x,int val)
{
while(x<n)
{
tree[x]+=val;
x+=~x&x+1;
}
}
int getsum(int x)
{
x=min(x,n-1);
int sum=0;
while(x>=0)
{
sum+=tree[x];
x-=~x&x+1;
}
return sum;
}
}bit[maxn<<1];
int fa[maxn];
int vis[maxn];
struct node
{
int v,next;
}edge[maxn*2];;
void init()
{
tot=0;
cnt=0;
tail=pool;
memset(head,-1,sizeof(head));
memset(vis,0,sizeof(vis));
}
void add_edge(int u,int v)
{
edge[tot].v=v;
edge[tot].next=head[u];
head[u]=tot++;
}
int dfssize(int u,int fa)
{
size[u]=1;
maxv[u]=0;
for(int i=head[u];i!=-1;i=edge[i].next)
{
int v=edge[i].v;
if(v==fa||vis[v])continue;
dfssize(v,u);
size[u]+=size[v];
if(size[v]>maxv[u])maxv[u]=size[v];
}
return size[u];
}
void dfsroot(int r,int u,int f)
{
if(size[r]-size[u]>maxv[u])
maxv[u]=size[r]-size[u];
if(maxv[u]<Max)Max=maxv[u],root=u;
for(int i=head[u];i!=-1;i=edge[i].next)
{
int v=edge[i].v;
if(v==f||vis[v])continue;
dfsroot(r,v,u);
}
}
int que[maxn];
int dis[maxn];
struct Node
{
int root,dis,subtree;
Node(){}
Node(int a,int b,int c):root(a),subtree(b),dis(c){}
};
vector<Node> V[maxn];
void solve(int u,int root,int subtree)
{
int l,r;
que[l=r=1]=u;
dis[u]=1;
fa[u]=0;
while(l<=r)
{
u=que[l];
V[u].push_back(Node(id[root],subtree,dis[u]));
for(int i=head[u];i!=-1;i=edge[i].next)
{
int v=edge[i].v;
if(vis[v]||v==fa[u])continue;
dis[v]=dis[u]+1;
que[++r]=v;
fa[v]=u;
}
l++;
}
bit[subtree].init(r+1);
for(int i=1;i<=r;i++)
{
int v=que[i];
bit[id[root]].add(dis[v],w[v]);
bit[subtree].add(dis[v],w[v]);
}
}
void dfs(int u)
{
Max=INF;
int tot=0;
tot=dfssize(u,0);
dfsroot(u,u,0);
vis[root]=1;
id[root]=cnt++;
V[root].push_back(Node(id[root],-1,0));
bit[id[root]].init(tot);
bit[id[root]].add(0,w[root]);
for(int i=head[root];i!=-1;i=edge[i].next)
{
int v=edge[i].v;
if(vis[v])continue;
solve(v,root,cnt);
cnt++;
}
for(int i=head[root];i!=-1;i=edge[i].next)
{
int v=edge[i].v;
if(vis[v])continue;
dfs(v);
}
}
int main()
{
while(scanf("%d%d",&N,&Q)!=EOF)
{
int u,v;
init();
for(int i=0;i<=N;i++)V[i].clear();
for(int i=1;i<=N;i++)scanf("%d",&w[i]);
for(int i=1;i<N;i++)
{
scanf("%d%d",&u,&v);
add_edge(u,v);
add_edge(v,u);
}
dfs(1);
char op[5];
int x,y;
while(Q--)
{
scanf("%s%d%d",op,&x,&y);
if(op[0]=='!')
{
int cha=y-w[x];
int len=V[x].size();
for(int i=0;i<len;i++)
{
Node tmp=V[x][i];
bit[tmp.root].add(tmp.dis,cha);
if(tmp.subtree!=-1)
bit[tmp.subtree].add(tmp.dis,cha);//?
}
w[x]+=cha;
}
else
{
int ans=0;
int len=V[x].size();
for(int i=0;i<len;i++)
{
Node tmp=V[x][i];
ans+=bit[tmp.root].getsum(y-tmp.dis);
if(tmp.subtree!=-1)
ans-=bit[tmp.subtree].getsum(y-tmp.dis);
}
printf("%d\n",ans);
}
}
}
return 0;
}