SPOJ 913 Query on a tree II 树链剖分

对于询问dist，树链剖分搞之，把边权转化到点上，然后注意细节就好（我在代码里标出来了，为了这个细节，wa了一屏）

对于询问kth，可以先求出两点(x和y)的lca，然后判断第k个数字是在x到lca的路径上还是y到lca的路径上，确定之后，倍增的寻找就好了~

View Code

#include <iostream>
#include <cstring>
#include <cstdlib>
#include <algorithm>
#include <cstdio>

#define N 50000
#define M 100000

using namespace std;

int head
,next[M],to[M],len[M];
int n,tot,cnt;
int fa
,son[M],top
,dat
,sum[N<<2],dep
,sz
,pre
,bh
;
int f
[22],bit[22];
int q[M];

inline void init()
{
memset(head,-1,sizeof head); cnt=2; tot=0;
memset(son,0,sizeof son);
memset(fa,0,sizeof fa);
memset(f,0,sizeof f);
memset(sum,0,sizeof sum);
bit[0]=1;
for(int i=1;i<=20;i++) bit[i]=bit[i-1]<<1;
}

inline void prep()
{
int h=1,t=2,sta;
q[1]=1; dep[1]=1;
while(h<t)
{
sta=q[h++]; sz[sta]=1;
for(int i=head[sta];~i;i=next[i])
if(fa[sta]!=to[i])
{
fa[to[i]]=sta;
f[to[i]][0]=sta;
pre[to[i]]=i^1;
dep[to[i]]=dep[sta]+1;
q[t++]=to[i];
}
}
for(int j=t-1;j>=1;j--)
{
sta=q[j];
for(int i=head[sta];~i;i=next[i])
if(fa[sta]!=to[i])
{
sz[sta]+=sz[to[i]];
if(sz[to[i]]>sz[son[sta]]) son[sta]=to[i];
}
}
for(int i=1;i<t;i++)
{
sta=q[i];
if(son[fa[sta]]==sta) top[sta]=top[fa[sta]];
else top[sta]=sta;
}
}

inline void rewrite()
{
for(int i=1;i<=n;i++)
if(top[i]==i)
for(int j=i;j;j=son[j])
{
bh[j]=++tot;
dat[tot]=len[pre[j]];
}
}

inline void lcainit()
{
for(int j=1;j<=20;j++)
for(int i=1;i<=n;i++)
f[i][j]=f[f[i][j-1]][j-1];
}

inline void pushup(int x)
{
sum[x]=sum[x<<1]+sum[x<<1|1];
}

inline void build(int u,int L,int R)
{
if(L==R) {sum[u]=dat[L];return;}
int MID=(L+R)>>1;
build(u<<1,L,MID); build(u<<1|1,MID+1,R);
pushup(u);
}

inline void add(int u,int v,int w)
{
to[cnt]=v; len[cnt]=w; next[cnt]=head[u]; head[u]=cnt++;
}

inline void read()
{
init();
scanf("%d",&n);
for(int i=1,a,b,c;i<n;i++)
{
scanf("%d%d%d",&a,&b,&c);
add(a,b,c); add(b,a,c);
}
prep();
rewrite();
build(1,1,tot);
lcainit();
}

inline int querysum(int u,int L,int R,int l,int r)
{
if(l<=L&&R<=r) return sum[u];
int MID=(L+R)>>1,res=0;
if(l<=MID) res+=querysum(u<<1,L,MID,l,r);
if(MID<r) res+=querysum(u<<1|1,MID+1,R,l,r);
return res;
}

inline int getsum(int x,int y)
{
int res=0;
while(top[x]!=top[y])
{
if(dep[top[x]]<dep[top[y]]) swap(x,y);
res+=querysum(1,1,tot,bh[top[x]],bh[x]);
x=fa[top[x]];
}
if(x==y) return res;//这句话好坑啊！把边权转移到点权上时会出现这个问题！
if(bh[x]>bh[y]) swap(x,y);
res+=querysum(1,1,tot,bh[son[x]],bh[y]);//细节
return res;
}

inline int getlca(int x,int y)
{
if(dep[x]<dep[y]) swap(x,y);
for(int i=20;i>=0;i--)
if(dep[f[x][i]]>=dep[y]) x=f[x][i];
if(x==y) return x;
for(int i=20;i>=0;i--)
if(f[x][i]!=f[y][i]) x=f[x][i],y=f[y][i];
return f[x][0];
}

inline int getlen(int x,int lca)
{
int res=0;
for(int i=20;i>=0;i--)
if(dep[f[x][i]]>=dep[lca]) x=f[x][i],res+=bit[i];
return res;
}

inline int getnum(int x,int p)
{
int res=0;
for(int i=20;i>=0;i--)
if(res+bit[i]<=p) x=f[x][i],res+=bit[i];
return x;
}

inline int getkth(int x,int y,int p)
{
int lca=getlca(x,y);
int lx=getlen(x,lca)+1;
int ly=getlen(y,lca)+1;
if(lx>=p) return getnum(x,p-1);
return getnum(y,lx+ly-p-1);
}

inline void go()
{
char str[10];int a,b,c;
while(scanf("%s",str))
{
if(str[1]=='O') break;
if(str[0]=='K')
{
scanf("%d%d%d",&a,&b,&c);
printf("%d\n",getkth(a,b,c));
}
else
{
scanf("%d%d",&a,&b);
printf("%d\n",getsum(a,b));
}
}
puts("");
}

int main()
{
int cas;scanf("%d",&cas);
while(cas--) read(),go();
return 0;
}