SPOJ QTREE2 - Query on a tree II [tree chain division + LCA]

Topic description

You are given a tree (an undirected acyclic connected graph) with N nodes, and edges numbered 1, 2, 3...N-1. Each edge has an integer value assigned to it, representing its length.

We will ask you to perfrom some instructions of the following form:

DIST a b : ask for the distance between node a and node b
or
KTH a b k : ask for the k-th node on the path from node a to node b

Given a tree of n points, edges have edge weights. The following actions are required:

DIST ab asks the sum of edge weights on the path from point a to point b

KTH abk Query the number of the kth point on the directed path from point a to point b

There are multiple sets of test data, and each set of data ends with DONE.

Example:

N = 6
1 2 1 // edge connects node 1 and node 2 has cost 1
2 4 1
2 5 2
1 3 1
3 6 2

Path from node 4 to node 6 is 4 -> 2 -> 1 -> 3 -> 6
DIST 4 6 : answer is 5 (1 + 1 + 1 + 2 = 5)
KTH 4 6 4 : answer is 3 (the 4-th node on the path from node 4 to node 6 is 3)

Input format:

The first line of input contains an integer t, the number of test cases (t <= 25). t test cases follow.

For each test case:

In the first line there is an integer N (N <= 10000)
In the next N-1 lines, the i-th line describes the i-th edge: a line with three integers a b c denotes an edge between a, b of cost c (c <= 100000)
The next lines contain instructions "DIST a b" or "KTH a b k"
The end of each test case is signified by the string "DONE".
There is one blank line between successive tests.

Output format:

For each "DIST" or "KTH" operation, write one integer representing its result.

Print one blank line after each test.
***********************************

Topic Analysis:

The main explanation is that kth asks
the path from u to v, which can be seen as u to lca(u,v) and lca(u,v) to v
. The kth node may be in the former or in the latter

When dep[u]-dep[lca]+1>=k , the kth node is in the former

while(dep[top[u]]>dep[lca])
{
    int dis=dep[u]-dep[top[u]]+1;//每次求出u到top[u]经过的节点数
    if(dis>=k) break;//如果经过节点数>=k说明带查询结点在这条中路径中
    k-=dis;//否则继续向上寻找
    u=fa[top[u]];
}
return pos[num[u]-k+1];

When dep[u]-dep[lca]+1< k , the kth node is in the latter

k-=dep[u]-dep[lca]+1;//先令k减去u到lca(u,v)经过的节点数
k=dep[v]-dep[lca]-k+1;//令k变成从v到lca(u,v)路径上的第k个
while(dep[top[v]]>dep[lca])//同上述
{
    int dis=dep[v]-dep[top[v]]+1;
    if(dis>=k) break;
    k-=dis;
    v=fa[top[v]];
}
return pos[num[v]-k+1];

#include<iostream>
#include<vector>
#include<algorithm>
#include<queue>
#include<cstring>
#include<cstdio>
using namespace std;
 
int read()
{
    int f=1,x=0;
    char ss=getchar();
    while(ss<'0'||ss>'9'){if(ss=='-')f=-1;ss=getchar();}
    while(ss>='0'&&ss<='9'){x=x*10+ss-'0';ss=getchar();}
    return f*x;
}

int t,n;
struct node{int v,dis,nxt;}E[30010];
int head[10010],tot;
int dep[10010],fa[10010],son[10010],size[10010];
int d[10010],top[10010],num[10010],pos[50010],cnt;
int sum[50010];
char ss[20];

void add(int u,int v,int dis)
{
    E[++tot].nxt=head[u];
    E[tot].v=v; E[tot].dis=dis;
    head[u]=tot;
} 

void dfs1(int u,int pa)
{
    size[u]=1;
    for(int i=head[u];i;i=E[i].nxt)
    {
        int v=E[i].v;
        if(v==pa) continue;
        dep[v]=dep[u]+1; fa[v]=u; d[v]=E[i].dis;
        dfs1(v,u);
        size[u]+=size[v];
        if(size[v]>size[son[u]]) son[u]=v;
    }
}

void dfs2(int u,int tp) 
{
    top[u]=tp; num[u]=++cnt; pos[cnt]=u;
    if(son[u]) dfs2(son[u],tp);
    for(int i=head[u];i;i=E[i].nxt)
    {
        int v=E[i].v;
        if(v==fa[u]||v==son[u]) continue;
        dfs2(v,v);
    }
}

void build(int s,int t,int p)
{
    if(s==t){sum[p]=d[pos[s]]; return;}
    int mid=s+t>>1;
    build(s,mid,p<<1); build(mid+1,t,p<<1|1);
    sum[p]=sum[p<<1]+sum[p<<1|1]; 
}

int get(int ll,int rr,int s,int t,int p)
{
    if(ll<=s&&t<=rr) return sum[p];
    int mid=s+t>>1;
    int ans=0;
    if(ll<=mid) ans+=get(ll,rr,s,mid,p<<1);
    if(rr>mid) ans+=get(ll,rr,mid+1,t,p<<1|1);
    return ans; 
}

int qsum()
{
    int u=read(),v=read();
    int ans=0;
    while(top[u]!=top[v])
    {
        if(dep[top[u]]<dep[top[v]]) swap(u,v);
        ans+=get(num[top[u]],num[u],1,n,1);
        u=fa[top[u]];
    }
    if(u==v) return ans;
    if(dep[u]>dep[v]) swap(u,v);
    return ans+=get(num[u]+1,num[v],1,n,1);
    
}

int getlca(int u,int v)
{
    while(top[u]!=top[v])
    {
        if(dep[top[u]]>dep[top[v]]) u=fa[top[u]];
        else v=fa[top[v]];
    }
    if(dep[u]<dep[v])return u;
    else return v;
}

int kth()
{
    int u=read(),v=read(),k=read();
    int lca=getlca(u,v);//求lca
    if(dep[u]-dep[lca]+1>=k)//k在u到lca(u,v)的路径中
    {
        while(dep[top[u]]>dep[lca])
        {
            int dis=dep[u]-dep[top[u]]+1;//每次求出u到top[u]经过的节点数
            if(dis>=k) break;//如果经过节点数>=k说明带查询结点在这条中路径中
            k-=dis;//否则继续向上寻找
            u=fa[top[u]];
        }
        return pos[num[u]-k+1];
    }
    else//k在lca(u,v)到v的路径中
    {
        k-=dep[u]-dep[lca]+1;//先令k减去u到lca(u,v)经过的节点数
        k=dep[v]-dep[lca]-k+1;//令k变成从v到lca(u,v)路径上的第k个
        while(dep[top[v]]>dep[lca])
        {
            int dis=dep[v]-dep[top[v]]+1;
            if(dis>=k) break;
            k-=dis;
            v=fa[top[v]];
        }
        return pos[num[v]-k+1];
    }
}

void init()
{
    cnt=tot=0;
    memset(head,0,sizeof(head));
    memset(fa,0,sizeof(fa));
    memset(son,0,sizeof(son));
    memset(top,0,sizeof(top));
    memset(num,0,sizeof(num));
    memset(d,0,sizeof(d));
    memset(sum,0,sizeof(sum));
}

int main()
{
    t=read();
    while(t--)
    {
        n=read();
        for(int i=1;i<n;++i)
        {
            int u=read(),v=read(),dis=read();
            add(u,v,dis);add(v,u,dis);
        }
    
        dep[1]=1;
        dfs1(1,0);dfs2(1,1);
        build(1,n,1);
        
        while(1)
        {
            scanf("%s",ss);
            if(ss[1]=='O') break;
            if(ss[1]=='I') printf("%d\n",qsum());
            else if(ss[1]=='T') printf("%d\n",kth());
        }
        init();
    }
    return 0;
}

Guess you like

Origin http://43.154.161.224:23101/article/api/json?id=324846055&siteId=291194637