SPOJ——QTREE2 - Query on a tree II（树链剖分或者倍增法）

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

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

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

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

Print one blank line after each test.

Example

Input:
1

6
1 2 1
2 4 1
2 5 2
1 3 1
3 6 2
DIST 4 6
KTH 4 6 4
DONE

Output:
5
3

题意：在树上查询两点之间距离以及路径上第k个点的编号

思路：看到树上操作直接想到树链剖分。。没想到用倍增法可以直接搞，也没有修改的操作。

所以我的代码会很冗长。。

#include <stdio.h>
#include <string.h>
#include <iostream>
#include <algorithm>
#include <vector>
#include <queue>
#include <set>
#include <map>
#include <string>
#include <math.h>
#include <stdlib.h>
using namespace std;

const int MAXN = 20010;
struct Edge
{
    int to,next;
}edge[MAXN*2];
int head[MAXN],tot;

int top[MAXN];
int fa[MAXN];
int deep[MAXN];
int num[MAXN];
int p[MAXN];
int fp[MAXN];
int son[MAXN];
int pos;

int n;

void init()
{
    tot = 0;
    memset(head,-1,sizeof(head));
    pos = 1;
    memset(son,-1,sizeof(son));
}
void addedge(int u,int v)
{
    edge[tot].to = v;edge[tot].next = head[u];head[u] = tot++;
}
void dfs1(int u,int pre,int d)
{
    deep[u] = d;
    fa[u] = pre;
    num[u] = 1;
    for(int i = head[u];i != -1; i = edge[i].next)
    {
        int v = edge[i].to;
        if(v != pre)
        {
            dfs1(v,u,d+1);
            num[u] += num[v];
            if(son[u] == -1 || num[v] > num[son[u]])
                son[u] = v;
        }
    }
}

void getpos(int u,int sp)
{
    top[u] = sp;
    if(son[u] != -1)
    {
        p[u] = pos++;
        fp[p[u]] = u;
        getpos(son[u],sp);
    }
    else
    {
        p[u] = pos++;
        fp[p[u]] = u;
        return;
    }
    for(int i = head[u] ; i != -1; i = edge[i].next)
    {
        int v = edge[i].to;
        if(v != son[u] && v != fa[u])
            getpos(v,v);
    }
}

#define lson l,m,rt<<1
#define rson m+1,r,rt<<1|1
int val[MAXN];
unsigned long long sum[MAXN<<2];

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

void build(int l,int r,int rt){
    if(l==r){
        sum[rt]=val[l];
        return;
    }
    int m=(l+r)>>1;
    build(lson);
    build(rson);
    pushup(rt);
}

void update(int p,int x,int l,int r,int rt){
    if(l==r){
        sum[rt]=x;
        return;
    }
    int m=(l+r)>>1;
    if(p<=m)
        update(p,x,lson);
    else
        update(p,x,rson);
    pushup(rt);
}



unsigned long long query_sum(int L,int R,int l,int r,int rt){
    if(l>=L&&R>=r){
        return sum[rt];
    }
    int m=(l+r)>>1;
    long long res=0;
    if(m>=L)
        res+=query_sum(L,R,lson);
    if(R>m)
        res+=query_sum(L,R,rson);
    return res;
}


unsigned long long _find_sum(int u,int v){
    int f1=top[u],f2=top[v];
    unsigned long long temp=0;
    while(f1!=f2){
        if(deep[f1]<deep[f2]){
            swap(f1,f2);
            swap(u,v);
        }
        temp+=query_sum(p[f1],p[u],1,n,1);
        u=fa[f1];f1=top[u];
    }
    if(u==v)
        return temp;
    if(deep[u]>deep[v])
        swap(u,v);
    temp+=query_sum(p[son[u]],p[v],1,n,1);
    return temp;
}

int _find(int u,int v,int k){
    int s1=0;
    int s2=0;
    int temp=v;
    if(k==1)
        return u;
    while(deep[u]>deep[v]){
        u=fa[u];
        s1++;
        if(s1==k-1)
            return u;
    }
    while(deep[temp]>deep[u]){
        temp=fa[temp];
        s2++;
    }
    while(u!=temp){
        u=fa[u];
        temp=fa[temp];
        s1++;
        s2++;
        if(s1==k-1)
            return u;
    }
    int cnt=0;
    if(cnt==s1-k+s2+1)
            return v;
    while(temp!=v){
        v=fa[v];
        cnt++;
        if(cnt==s1-k+s2+1)
            return v;
    }
}

int e[MAXN][3];
int main(){
    int t;
    scanf("%d",&t);
    while(t--){
        init();
        scanf("%d",&n);
        getchar();
        for(int i=0;i<n-1;i++){
            scanf("%d%d%d",&e[i][0],&e[i][1],&e[i][2]);
            addedge(e[i][0],e[i][1]);
            addedge(e[i][1],e[i][0]);
        }
        dfs1(1,0,0);
        getpos(1,1);
        for(int i=0;i<n-1;i++){
            if(deep[e[i][0]]<deep[e[i][1]])
                swap(e[i][0],e[i][1]);
            val[p[e[i][0]]]=e[i][2];
        }
        build(1,n,1);
        char op[10];
        int u,v,k;


        while(scanf("%s",op)){
            if(op[1]=='O'){
                printf("\n");
                break;
            }
            if(op[1]=='I'){
                scanf("%d%d",&u,&v);
                printf("%llu\n",_find_sum(u,v));
            }
            else{
                scanf("%d%d%d",&u,&v,&k);
                printf("%d\n",_find(u,v,k));
            }
        }
    }
    return 0;
}