loj#2322. 「清华集训 2017」Hello world!【分块+并查集+树状数组】

传送门

解题思路：

考虑对步数分块，对于步数大于S的直接暴力修改查询，步数小于S的建S棵树，第i棵树x先第i个祖先连边，用并查集维护连续是1的一段，树状数组维护查询。
这样修改是$O(Slogn)$的，查询是$O(logn)$，不知有没有更优秀的复杂度。

#include<iostream>
#include<algorithm>
#include<ctime>
#define ll long long
#define pb push_back
using namespace std;
ll getint()
{
    ll i=0,f=1;char c;
    for(c=getchar();(c!='-')&&(c<'0'||c>'9');c=getchar());
    if(c=='-')c=getchar(),f=-1;
    for(;c>='0'&&c<='9';c=getchar())i=(i<<3)+(i<<1)+c-'0';
    return i*f;
}
const int N=50005,B=20;
int n,dep[N],fa[N][20];ll a[N];
vector<int>g[N];
void dfs(int u)
{
    for(int i=1;i<17;i++)fa[u][i]=fa[fa[u][i-1]][i-1];
    for(int v:g[u])if(v!=fa[u][0])
        fa[v][0]=u,dep[v]=dep[u]+1,dfs(v);
}
int jump(int x,int det)
{
    for(int i=0;(1<<i)<=det;i++)
        if(det>>i&1)x=fa[x][i];
    return x;
}
int lca(int x,int y)
{
    if(dep[x]<dep[y])swap(x,y);
    x=jump(x,dep[x]-dep[y]);if(x==y)return x;
    for(int i=16;i>=0;i--)if(fa[x][i]!=fa[y][i])x=fa[x][i],y=fa[y][i];
    return fa[x][0];
}
struct Tree
{
    int idx,fa[N],top[N],st[N],ed[N];ll bit[N];
    vector<int>g[N];
    void add(int i,ll v){for(;i<=n;i+=i&-i)bit[i]+=v;}
    ll query(int i,ll res=0){for(;i>0;i-=i&-i)res+=bit[i];return res;}
    int find(int x){return top[x]==x?x:top[x]=find(top[x]);}
    void dfs(int u)
    {
        if(u)
        {
            st[u]=++idx;
            if(a[u]==1)top[u]=find(fa[u]);
        }
        for(int v:g[u])dfs(v);
        if(!u)return;
        ed[u]=idx;
        add(st[u],a[u]),add(ed[u]+1,-a[u]);
    }
    void init(int S)
    {
        for(int i=1;i<=n;i++)
        {
            top[i]=i;fa[i]=jump(i,S);
            g[fa[i]].pb(i);
        }
        dfs(0);
    }
    void modify(int u,ll x,ll y)
    {
        add(st[u],y-x),add(ed[u]+1,x-y);
        if(y==1)top[u]=find(fa[u]);
    }
    ll query(int x,int y)
    {
        return query(st[x])-query(st[y]);
    }
}T[B+1];
void update(int x)
{
    if(a[x]==1||!x)return;
    ll t=sqrt(a[x]);
    for(int i=1;i<=B;i++)T[i].modify(x,a[x],t);
    a[x]=t;
}
void Modify(int x,int y,int k)
{
    int t=lca(x,y),dis=dep[x]+dep[y]-2*dep[t],det=dis%k;
    if(det)
    {
        update(y);
        if(dep[y]-dep[t]<det)y=t=jump(x,dis-det);
        else y=jump(y,det);
    }
    if(k>B)
    {
        while(dep[x]>=dep[t])update(x),x=jump(x,k);
        while(dep[y]>dep[t])update(y),y=jump(y,k);
    }
    else
    {
        while(dep[x]>=dep[t])update(x),x=T[k].fa[x],x=T[k].find(x);
        while(dep[y]>dep[t])update(y),y=T[k].fa[y],y=T[k].find(y);
    }
}
ll Query(int x,int y,int k)
{
    ll res=0;
    int t=lca(x,y),dis=dep[x]+dep[y]-2*dep[t],det=dis%k;
    if(det)
    {
        res+=a[y];
        if(dep[y]-dep[t]<det)y=t=jump(x,dis-det);
        else y=jump(y,det);
    }
    if(k>B)
    {
        while(dep[x]>=dep[t])res+=a[x],x=jump(x,k);
        while(dep[y]>dep[t])res+=a[y],y=jump(y,k);
    }
    else
    {
        dis=dep[x]-dep[t];
        int z=T[k].fa[jump(x,dis-dis%k)];
        res+=T[k].query(x,z);
        dis=dep[y]-dep[t];
        z=dis%k?T[k].fa[jump(y,dis-dis%k)]:t;
        res+=T[k].query(y,z);
    }
    return res;
}
int main()
{
    //freopen("lx.in","r",stdin);
    //freopen("lx.out","w",stdout);
    n=getint();int op,x,y,k;
    for(int i=1;i<=n;i++)a[i]=getint();
    for(int i=1;i<n;i++)x=getint(),y=getint(),g[x].pb(y),g[y].pb(x);
    dep[1]=1,dfs(1);
    for(int i=1;i<=B;i++)T[i].init(i);
    for(int Q=getint();Q;Q--)
    {
        op=getint(),x=getint(),y=getint(),k=getint();
        if(!op)Modify(x,y,k);
        else printf("%lld\n",Query(x,y,k));
    }
    return 0;
}