HDU4918 Query on the subtree 点分治+树状数组

bobo has a tree, whose vertices are conveniently labeled by 1,2,…,n. At the very begining, the i-th vertex is assigned with weight w _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.

InputThe input consists of several tests. For each tests:

The first line contains n,q (1≤n,q≤10 ⁵). The second line contains n integers w ₁,w ₂,…,w _n (0≤w _i≤10 ⁴). 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 ⁴,0≤d≤n).
OutputFor each tests:

For each queries, a single number denotes the total weight.Sample Input

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

Sample Output

3
2
1
6
6

---

题意：给你一棵树，N个点，每个点一个权值，然后Q组操作（共两种），第一种是求导一个节点距离不超过d的所有点的权值和是多少；第二种操作时，修改一个点的权值；
题解：多次进行操作1。此时不能每次都O（NlogN）了，太慢了。我们考虑到对于点分治，树的重心一共有logN层，第一层为整棵树的重心，第二层为第一层重心的子树的重心，以此类推，每次至少分成两个大小差不多的子树，所以
一共有logN层。而且，对于一个点，他最多只属于logN个子树，也就是最多只属于logN个重心。所以我们可以预处理出每个点所属于的重心以及到这些重心的距离，以每个重心建树状数组，每个点按照到重心的距离插入到树状数组中，
然后每次查询到u距离不超过d的点的个数就通过树状数组求前缀和得到。假设一个重心x到u的距离为dis，那么便统计到重心x距离不超过d-dis的点的个数，这个过程我们称之为“借力”，本身能力有限，所以需要借助x的影响力。因为
如果这个重心被u借力了，那么这个重心的子重心一定也被借力，由于相邻被借力的两个重心x、y所统计的点会有重复，所以我们需要去重。去重的话我们就通过对每个节点再开一个v对x的树状数组，这个树状数组的意义为：重心x的子
树v的重心为y时，子树v中每个点到x的距离为下标建立的树状数组。因为重心x与重心y交集的部分，重心x包括的部分重心y一定包括，所以统计的时候减去v对x的树状数组中距x不超过d-dis的点的个数即可。访问u所属与的所有重心，
挨个借力，同时去重，便能得到距离u不超过d的点的个数。因为重心最多logN层，每个树状数组最多N个点，logN复杂度的统计，所以每次查询复杂度O（logN*logN）。我们最多为每个节点开2个树状数组，而且每一层所有树状数组的
大小相加不超过N，所以树状数组的占用空间为O（2NlogN）。
在上面的基础上稍做扩充。预处理的时候插入树状数组的就是该点的权值，查询依旧是统计前缀和。修改点权值的时候，便是和查询一样，在u距重心x距离d的位置在x的树状数组中修改u的权值，同时修改u属于重心x的子树v的v对x的树
状数组中相同位置的值。复杂度和查询一样为O（logN*logN）。

参考代码：

#include<bits/stdc++.h>
using namespace std;
#define pii pair<int,int>
#define mkp make_pair
#define lowbit(x) (x&-x)
 
typedef long long ll;
const int INF=0x3f3f3f3f;
const int maxn=1e5+;
int n,q,w[maxn];
char op[];
struct MSG{
    int id1,id2;
    int dep;
} msg[maxn][];
struct Edge{
    int v,nxt;
} edge[maxn<<];
int vis[maxn],head[maxn],tot;
int root,siz[maxn],mx[maxn],fa[maxn],S;
int maxfloor[maxn],idn,L[maxn<<],R[maxn<<];
int c[maxn*];
void Init()
{
    S=n;idn=;
    tot=root=;
    memset(c,,sizeof c);
    memset(vis,,sizeof vis);
    memset(w,,sizeof w);
    memset(head,-,sizeof head);
    memset(msg,,sizeof msg);
}
void AddEdge(int x,int y)
{
    edge[tot].v=y;
    edge[tot].nxt=head[x];
    head[x]=tot++;
}
 
void Add(int l,int r,int pos,int val)
{
    while(pos<r-l)
        c[l+pos]+=val,pos+=lowbit(pos);
}
int Sum(int l,int r,int len)
{
    if(len<) return ;
    if(len>r-l-) len=r-l-;
    int res=;
    while(len) res+=c[l+len],len-=lowbit(len);
    return res;
}
 
void getroot(int u,int fa)
{
    siz[u]=;mx[u]=;
    for(int i=head[u];~i;i=edge[i].nxt)
    {
        int v=edge[i].v;
        if(vis[v]||v==fa) continue;
        getroot(v,u);
        siz[u]+=siz[v];
        mx[u]=max(mx[u],siz[v]);
    }
    mx[u]=max(mx[u],S-siz[u]);
    if(mx[u]<mx[root]) root=u;
}
 
int getmaxdep(int u,int fa)
{
    int res=;
    for(int i=head[u];~i;i=edge[i].nxt)
    {
        int v=edge[i].v;
        if(vis[v]||v==fa) continue;
        res=max(res,+getmaxdep(v,u));
    }
    return res;
}
void dfs(int u,int fa,int deep,int id,int flor,int tp)
{
    if(!tp) msg[u][flor].id1=id;
    else msg[u][flor].id2=id;
    msg[u][flor].dep=deep;
    Add(L[idn],R[idn],deep,w[u]);
    for(int i=head[u];~i;i=edge[i].nxt)
    {
        int v=edge[i].v;
        if(!vis[v]&&v!=fa) dfs(v,u,deep+,id,flor,tp);
    }
}
void solve(int u,int s,int flor)
{
    vis[u]=; maxfloor[u]=flor;
    idn++; L[idn]=R[idn-];
    R[idn]=L[idn]+getmaxdep(u,)+;
    dfs(u,,,idn,flor,);
    msg[u][flor].id2=-;
    for(int i=head[u];~i;i=edge[i].nxt)
    {
        int v=edge[i].v;
        if(vis[v]) continue;
        idn++;L[idn]=R[idn-];
        R[idn]=L[idn]+getmaxdep(v,u)+;
        dfs(v,u,,idn,flor,);
    }
 
     for(int i=head[u];~i;i=edge[i].nxt)
    {
        int v=edge[i].v;
        if(vis[v]) continue;
        S=siz[v]; root=;
        if(siz[v]>siz[u]) S=s-siz[u];
        getroot(v,u);
        solve(root,siz[v],flor+);
    }
}
 
int main()
{
    while(~scanf("%d%d",&n,&q))
    {
        Init();
        for(int i=;i<=n;++i) scanf("%d",w+i);
        for(int i=;i<n;++i)
        {
            int x,y;
            scanf("%d%d",&x,&y);
            AddEdge(x,y);AddEdge(y,x);
        }
        mx[root]=INF;
        getroot(,);
        solve(,S,);
        while(q--)
        {
            int x,y,ans;
            scanf("%s%d%d",&op,&x,&y);
            if(op[]=='?')
            {
                ans=;
                for(int f=;f<=maxfloor[x];++f)
                {
                    int id1=msg[x][f].id1;
                    int id2=msg[x][f].id2;
                    int dep=msg[x][f].dep;
                    ans+=Sum(L[id1],R[id1],y+-dep);
                    if(id2!=-) ans-=Sum(L[id2],R[id2],y+-dep);
                }
                printf("%d\n",ans);
            }
            else
            {
                for(int f=;f<=maxfloor[x];++f)
                {
                    int id1=msg[x][f].id1;
                    int id2=msg[x][f].id2;
                    int dep=msg[x][f].dep;
                    Add(L[id1],R[id1],dep,y-w[x]);
                    if(id2!=-) Add(L[id2],R[id2],dep,y-w[x]);
                }
                w[x]=y;
            }
        }
    }
 
    return ;
}