SP913 QTREE2 - Query on a tree II

思路

第一个可以倍增，第二个讨论在a到lca的路径上还是lca到b的路径上，
倍增即可

代码

#include <cstdio>
#include <algorithm>
#include <cstring>
using namespace std;
int jump[10010][16],sum[10010][16],fir[10010],nxt[10010*2],v[10010*2],w[10010*2],cnt,dep[10010],n;
void addedge(int ui,int vi,int wi){
    ++cnt;
    v[cnt]=vi;
    w[cnt]=wi;
    nxt[cnt]=fir[ui];
    fir[ui]=cnt;
}
void dfs(int u,int f,int wx){
    dep[u]=dep[f]+1;
    jump[u][0]=f;
    sum[u][0]=wx;
    for(int i=1;i<15;i++)
        jump[u][i]=jump[jump[u][i-1]][i-1],sum[u][i]=sum[u][i-1]+sum[jump[u][i-1]][i-1];
    for(int i=fir[u];i;i=nxt[i]){
        if(v[i]==f)
            continue;
        dfs(v[i],u,w[i]);
    }
}
int lca(int x,int y){
    if(dep[x]<dep[y])
        swap(x,y);
    for(int i=15;i>=0;i--)
        if(dep[jump[x][i]]>=dep[y])
            x=jump[x][i];
    if(x==y)
        return x;
    for(int i=15;i>=0;i--)
        if(jump[x][i]!=jump[y][i])
            x=jump[x][i],y=jump[y][i];
    return jump[x][0];
}
int query_sum(int x,int y){
    int ans=0;
    if(dep[x]<dep[y])
        swap(x,y);
    for(int i=15;i>=0;i--)
        if(dep[jump[x][i]]>=dep[y]){
            ans+=sum[x][i];
            x=jump[x][i];
        }
    if(x==y)
        return ans;
    for(int i=15;i>=0;i--)
        if(jump[x][i]!=jump[y][i]){
            ans+=sum[x][i]+sum[y][i];
            x=jump[x][i],y=jump[y][i];
        }
    return ans+sum[x][0]+sum[y][0];
}
int kth_fa(int x,int k){
    for(int i=15;i>=0;i--){
        if((k>>i)&1)
            x=jump[x][i];
    }
    return x;
}
void init(void){
    memset(jump,0,sizeof(jump));
    memset(sum,0,sizeof(sum));
    memset(fir,0,sizeof(fir));
    memset(nxt,0,sizeof(nxt));
    memset(v,0,sizeof(v));
    memset(w,0,sizeof(w));
    cnt=0;
    memset(dep,0,sizeof(dep));
}
int main(){
    int T;
    scanf("%d",&T);
    while(T--){
        init();
        scanf("%d",&n);
        for(int i=1;i<n;i++){
            int a,b,c;
            scanf("%d %d %d",&a,&b,&c);
            addedge(a,b,c);
            addedge(b,a,c);
        }
        dfs(1,0,0);
        char opt[10];
        while(1){
            scanf("%s",opt);
            if(opt[1]=='O')
                break;
            if(opt[1]=='I'){
                int a,b;
                scanf("%d %d",&a,&b);
                printf("%d\n",query_sum(a,b));
            }
            else{
                int a,b,c;
                scanf("%d %d %d",&a,&b,&c);
                int LCA=lca(a,b);
                if(dep[a]-dep[LCA]+1>=c){
                    printf("%d\n",kth_fa(a,c-1));
                }
                else{
                    c-=dep[a]-dep[LCA];
                    int t2=dep[b]-dep[LCA]+1;
                    printf("%d\n",kth_fa(b,t2-c));
                }
            }
        }
    }
    return 0;
}