SPOJ QTREE3 Query on a tree again! (树链剖分)

题目链接：https://vjudge.net/contest/314922#problem/B

#include <cmath>
#include <cstdio>
#include <algorithm>
#include <iostream>
#include <cstring>
#include <ctype.h>
#include <map>
#include <vector>
#include <set>
#include <bitset>
#define lson rt << 1, l, mid
#define rson rt << 1 | 1, mid + 1, r
#define lc rt << 1
#define rc rt << 1 | 1
#define INF 0x7fffffff
#define eps 1e-9
#define pi acos(-1.0)
using namespace std;
typedef long long ll;
static const int MAX_N = 1e5 + 5;
static const int Mod = 10007;
struct Edge{
    int to, next;
}edge[MAX_N << 1];
int head[MAX_N];    //链式前向星
int siz[MAX_N], son[MAX_N], top[MAX_N], fa[MAX_N], tid[MAX_N], dep[MAX_N], pos[MAX_N];
int cnt, tot;
void dfs1(int u, int f, int d){     //第一次dfs得到重儿子
    siz[u] = 1;
    dep[u] = d;
    son[u] = 0;
    fa[u] = f;
    for(int i = head[u]; i != -1; i = edge[i].next){
        int v = edge[i].to;
        if(v == f) continue;    //不走重复路
        dfs1(v, u, d + 1);
        siz[u] += siz[v];
        if(siz[v] > siz[son[u]]) son[u] = v;
    }
}
void dfs2(int u,int t){     //将重儿子连接成重链,轻儿子连接成轻链(得到区间)
    top[u] = t;
    tid[u] = ++cnt;
    pos[cnt] = u;
    if(son[u]) dfs2(son[u],t);  //优先处理重链
    for(int i = head[u]; i != -1; i = edge[i].next){
        int v = edge[i].to;
        if(v != son[u] && v != fa[u]) dfs2(v,v);
    }
}
int tag[MAX_N << 2];
void prework(){
    memset(head, -1, sizeof(head));
    memset(tag, 0, sizeof(tag));
    cnt = 0;
    tot = 0;
}
void addEdge(int u, int v){
    edge[tot].to = v;
    edge[tot].next = head[u];
    head[u] = tot++;
}
void pushUp(int rt){
    tag[rt] = tag[lc] + tag[rc];
}
void update(int rt, int l, int r, int pos){
    if(l == r){
        tag[rt] ^= 1;
        return ;
    }
    int mid = l + r >> 1;
    if(pos <= mid) update(lson, pos);
    else update(rson, pos);
    pushUp(rt);
}
int queryInd(int rt, int l, int r, int ql, int qr){
    if(!tag[rt]) return -1;
    if(l == r) return pos[l];
    int mid = l + r >> 1;
    int ans = -1;
    if(ql <= mid) ans = queryInd(lson, ql, qr);
    if(ans != -1) return ans;
    if(qr > mid) ans = queryInd(rson, ql, qr);
    return ans;
}
int queryPathInd(int u, int v){
    int ans = -1, temp;
    while(top[u] != top[v]){
        if(dep[top[u]] < dep[top[v]]) swap(u, v);
        temp = queryInd(1, 1, cnt, tid[top[u]], tid[u]);
        if(temp != -1) ans = temp;
        u = fa[top[u]];
    }
    if(dep[u] > dep[v]) swap(u, v);
    temp = queryInd(1, 1, cnt, tid[u], tid[v]);  //后面离的越近
    if(temp != -1) ans = temp;
    return ans;
}
void mainwork(){
    int n, q;
    scanf("%d%d", &n, &q);
    for(int i = 1; i < n; ++i){
        int u, v;
        scanf("%d%d", &u, &v);
        addEdge(u, v);
        addEdge(v, u);
    }
    dfs1(1, 0, 1);
    dfs2(1, 1);
    while(q--){
        int opt, x;
        scanf("%d%d", &opt, &x);
        if(opt) printf("%d\n", queryPathInd(1, x));
        else update(1, 1, cnt, tid[x]);
    }
}
int main(){
//    freopen("input.txt", "r", stdin);
//    freopen("output.txt", "w", stdout);
    prework();
    mainwork();
    return 0;
}

View Code