bzoj 4381: [POI2015]Odwiedziny 分块+树链剖分

题意

给定一棵n个点的树，树上每条边的长度都为1，第i个点的权值为a[i]。
Byteasar想要走遍这整棵树，他会按照某个1到n的全排列b走n-1次，第i次他会从b[i]点走到b[i+1]点，并且这一次的步伐大小为c[i]。
对于一次行走，假设起点为x，终点为y，步伐为k，那么Byteasar会从x开始，每步往前走k步，如果最后不足k步就能到达y，那么他会一步走到y。
请帮助Byteasar统计出每一次行走时经过的所有点的权值和。
2<=n<=50000，1<=a[i]<=10000

分析

看到n的大小就不难想到要用到分块。于是我们可以按步伐大于根号n和小于根号n将询问分为两类。对于大于根号n的部分，我们只要每次暴力跳即可。小于根号n的部分就预处理。
具体实现起来就是，预处理up[x,y]表示从点x开始每次跳y步直到不能跳为止的权值和，然后便可以O(logn)处理第二部分的询问。
对于暴力跳而言，也就是要支持快速查找第k级祖先，这有一个经典的长链剖分做法，但放在这题中是可以用重链剖分来实现的。在找一个点的k级祖先时，若k级祖先在当前重链上，就根据其在dfs序中的位置直接得到祖先，否则就暴力往上跳。那么处理询问的总复杂度就是O(logn+sqrtn)。

代码

#include<iostream>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<algorithm>
#include<cmath>
using namespace std;

const int N=50005;

int n,a[N],b[N],c[N],cnt,last[N],dfn[N],pos[N],tim,top[N],fa[N],size[N],dep[N],up[N][235],block;
struct edge{int to,next;}e[N*2];

int read()
{
    int x=0,f=1;char ch=getchar();
    while (ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
    while (ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
    return x*f;
}

void addedge(int u,int v)
{
    e[++cnt].to=v;e[cnt].next=last[u];last[u]=cnt;
    e[++cnt].to=u;e[cnt].next=last[v];last[v]=cnt;
}

void dfs1(int x)
{
    dep[x]=dep[fa[x]]+1;size[x]=1;
    int y=x;
    for (int i=1;i<=block;i++) y=fa[y],up[x][i]=a[x]+up[y][i];
    for (int i=last[x];i;i=e[i].next)
    {
        if (e[i].to==fa[x]) continue;
        fa[e[i].to]=x;
        dfs1(e[i].to);
        size[x]+=size[e[i].to];
    }
}

void dfs2(int x,int chain)
{
    dfn[++tim]=x;pos[x]=tim;top[x]=chain;int k=0;
    for (int i=last[x];i;i=e[i].next)
        if (e[i].to!=fa[x]&&size[e[i].to]>size[k]) k=e[i].to;
    if (!k) return;
    dfs2(k,chain);
    for (int i=last[x];i;i=e[i].next)
        if (e[i].to!=fa[x]&&e[i].to!=k) dfs2(e[i].to,e[i].to);
}

int get_lca(int x,int y)
{
    while (top[x]!=top[y])
    {
        if (dep[top[x]]<dep[top[y]]) swap(x,y);
        x=fa[top[x]];
    }
    return dep[x]<dep[y]?x:y;
}

int get_kth(int x,int k)
{
    while (x&&dep[x]-dep[top[x]]<k) k-=dep[x]-dep[top[x]]+1,x=fa[top[x]];
    return !x?0:dfn[pos[x]-k];
}

int main()
{
    n=read();block=sqrt(n);
    for (int i=1;i<=n;i++) a[i]=read();
    for (int i=1;i<n;i++)
    {
        int x=read(),y=read();
        addedge(x,y);
    }
    for (int i=1;i<=n;i++) b[i]=read();
    for (int i=1;i<n;i++) c[i]=read();
    dfs1(1);dfs2(1,1);
    for (int i=1;i<n;i++)
    {
        int x=b[i],y=b[i+1],z=c[i],lca=get_lca(x,y);
        if (z<=block)
        {
            int ans=up[x][z]-up[get_kth(lca,z-(dep[x]-dep[lca])%z)][z];
            int tmp=get_kth(y,(dep[x]+dep[y]-dep[lca]*2)%z);
            if (dep[tmp]>=dep[lca]) ans+=up[tmp][z]-up[get_kth(lca,z-(dep[tmp]-dep[lca])%z)][z];
            if ((dep[x]+dep[y]-dep[lca]*2)%z>0) ans+=a[y];
            if ((dep[x]-dep[lca])%z==0) ans-=a[lca];
            printf("%d\n",ans);
        }
        else
        {
            int ans=0,tmp=x;
            while (dep[tmp]>=dep[lca]) ans+=a[tmp],tmp=get_kth(tmp,z);
            tmp=get_kth(y,(dep[x]+dep[y]-dep[lca]*2)%z);
            while (dep[tmp]>=dep[lca]) ans+=a[tmp],tmp=get_kth(tmp,z);
            if ((dep[x]+dep[y]-dep[lca]*2)%z>0) ans+=a[y];
            if ((dep[x]-dep[lca])%z==0) ans-=a[lca];
            printf("%d\n",ans);
        }
    }
    return 0;
}