SP913 QTREE2 - Query on a tree II

給定一棵n個點的樹，邊具有邊權。要求作以下操作：

DIST a b 詢問點a至點b路徑上的邊權之和

KTH a b k 詢問點a至點b有向路徑上的第k個點的編號

有多組測試數據，每組數據以DONE結尾。

裸的LCA。

在處理第二個操作時，我直接向上數跳了多少個。

顧z大佬說不能這麽做，要求出跳到那個點的深度再去跳。

真的是這樣，不過懶得想了，應該是+1-1的誤差。 balabala。。。

code:

#include <iostream>
#include <cstdio>
#include <cstring>

using namespace std;

const int wx=50017;

int dep[wx],dis[wx];
int f[wx][23];
int head[wx];
int num,n,t;
char opt[7];

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

struct e{
    int nxt,to,dis;
}edge[wx*2];

void add(int from,int to,int dis){
    edge[++num].nxt=head[from];
    edge[num].to=to;
    edge[num].dis=dis;
    head[from]=num;
}

void dfs(int u,int fa){
    dep[u]=dep[fa]+1;
    for(int i=head[u];i;i=edge[i].nxt){
        int v=edge[i].to;
        if(v==fa)continue;
        f[v][0]=u;dis[v]=dis[u]+edge[i].dis;
        dfs(v,u);
    }
}

void pre(){
    for(int j=1;j<=21;j++)for(int i=1;i<=n;i++)f[i][j]=f[f[i][j-1]][j-1];
}

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

int find(int x,int k){
    for(int i=21;i>=0;i--){
        if(dep[f[x][i]]>=k)x=f[x][i];
    }
    return x;
}

int main(){
    t=read();
    while(t--){
        n=read();
        memset(head,0,sizeof head); num=1;
        memset(edge,0,sizeof edge);
        for(int i=1;i<n;i++){
            int x,y,z;
            x=read(); y=read(); z=read();
            add(x,y,z); add(y,x,z);
        }
        dfs(1,0); pre();
        while(1){
            scanf("%s",opt+1);
            if(opt[2]=='O')break;
            if(opt[2]=='I'){
                int x,y;
                x=read(); y=read();
                int lca=LCA(x,y);
                printf("%d\n",dis[x]+dis[y]-2*dis[lca]);
            }
            if(opt[1]=='K'){
                int a,b,k;
                a=read(); b=read(); k=read();
                int lca=LCA(a,b);
                if(dep[a]-dep[lca]+1>=k)printf("%d\n",find(a,dep[a]-k+1));
                else printf("%d\n",find(b,k-dep[a]+2*dep[lca]-1));
            }
        }
    }
}
 

LCA SP913 QTREE2 - Query on a tree II