Query on a tree II LCA题 [SPOJ-QTREE2]

Posted 2021-03-12 xxfy1

其实用不上树链剖分也行，就普通的lca题目
但是经过这道题可以看出树链剖分也可以解决LCA的题目，
注意ptn是dfn的反查，这样能够快速得到dfs序中相应的节点

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int,int> pii;
const int N = 2e5;

int fa[N],lv[N],siz[N],son[N],top[N],dfn[N],cnt=0;
int ptn[N];
vector<pii> G[N];
int n,seg[N<<2];

int query(int s,int t,int l,int r,int p) {
    if(s<=l&&r<=t) return seg[p];
    int m=(l+r)/2;
    if(t<=m) return query(s,t,l,m,p<<1);
    if(m+1<=s) return query(s,t,m+1,r,p<<1|1);
    return query(s,t,l,m,p<<1)+query(s,t,m+1,r,p<<1|1);
}

void update(int pos,int val,int l,int r,int p){
    if(l==r) {
        seg[p]=val;
        return;
    }
    int m=(l+r)/2;
    if(pos<=m) update(pos,val,l,m,p<<1);
    else update(pos,val,m+1,r,p<<1|1);
    seg[p]=seg[p<<1]+seg[p<<1|1];
}

int dfs1(int x,int fx) {
    son[x]=-1;
    siz[x]=1;
    fa[x]=fx;
    lv[x]=lv[fx]+1;
    for(auto p:G[x]) {
        int to = p.first;
        if(to==fx) continue;
        siz[x] += dfs1(to,x);
        if(son[x] == -1 || siz[to] > siz[son[x]]) son[x] = to;
    }
    return siz[x];
}

void dfs2(int x,int root) {
    top[x]=root;
    dfn[x]=++cnt;
    ptn[dfn[x]]=x;
    if(son[x]==-1) return;
    dfs2(son[x],root);
    for(auto p:G[x]) {
        int to = p.first;
        if(to == son[x] || to == fa[x]) continue;
        dfs2(to,to);
    }
}

int dist(int x,int y) {
    int ans = 0;
    while(top[x] != top[y]) {
        if(lv[top[x]]<lv[top[y]]) swap(x,y);
        ans+=query(dfn[top[x]],dfn[x],1,n,1);
        x=fa[top[x]];
    }
    if(dfn[x] != dfn[y]) {
        if(dfn[x]>dfn[y]) swap(x,y);
        ans+=query(dfn[son[x]],dfn[y],1,n,1);
    }
    return ans;
}

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

int kth(int x,int y,int k) {
    int lca=querylca(x,y);
    if(lv[x]-lv[lca]+1>=k ) {
        while(lv[top[x]] > lv[lca]) {
            if(lv[x] - lv[top[x]] +1 >= k) break;
            k-=lv[x] - lv[top[x]] +1;
            x=fa[top[x]];
        }
        return ptn[dfn[x] -k+1 ];
    }
    else {
        k -= lv[x] - lv[lca] + 1;
        k = lv[y] - lv[lca] - k + 1;
        while(lv[top[y]] > lv[lca]) {
            if(lv[y] - lv[top[y]] + 1 >= k) break;
            k -= lv[y] - lv[top[y]] + 1;
            y = fa[top[y]];
        }
        return ptn[dfn[y] - k + 1];
    }
}

int a,b,c,k;

void solve() {
    memset(seg,0,sizeof(seg));
    cnt = 0;
//    cin>>n;
    scanf("%d",&n);
    for(int i=1;i<=n;i++) G[i].clear();
    for(int i=1;i<n;i++){
        scanf("%d %d %d",&a,&b,&c);
        G[a].push_back({b,c});
        G[b].push_back({a,c});
    }
    dfs1(1,-1);
    dfs2(1,1);
    for(int i=1;i<=n;i++){
        for(auto e:G[i]) {
            int to=e.first,cost=e.second;
            if(lv[to]<lv[i]) to=i;
            update(dfn[to],cost,1,n,1);
        }
    }
    char str[10];

//    for(int i=1;i<=n;i++) cout<<query(i,i,1,n,1)<<endl;
    while(true) {
        scanf(" %s",str);
        if(str[1] == 'I' ) {
//            cin>>a>>b;
            scanf("%d %d",&a,&b);
//            cout<<ask(a,b)<<'
';
            printf("%d
",dist(a,b));
        } else if(str[0] == 'K') {
//            cin>>a>>b;
            scanf("%d %d %d",&a,&b,&k);
            printf("%d
",kth(a,b,k));
        } else break;
    }
//    for(int i=1;i<=n;i++) cout<<query(i,i,1,n,1)<<endl;
}

int main(){
    ios::sync_with_stdio(false);
    int t;
    scanf("%d",&t);
    while(t--) solve();
}