(SPOJ) QTREE - Query on a tree [樹鍊剖分]

http://www.spoj.com/problems/QTREE/

樹鍊剖分

不過在TOJ上面的卻TLE了QQ

#include <iostream>
#include <stdio.h>
#include <vector>
#include <cstring>
#include <algorithm>
using namespace std;

const int MAX_N = 1e5 +6;
const int MAX_P = 17;

vector<int> edg[MAX_N];
int seg[4*MAX_N];
int val[MAX_N];

void init(int id,int L,int R) {
    if (L==R) {
        seg[id] = val[L];
        return;
    }
    int mid=(L+R)>>1;
    init(id*2,L,mid);
    init(id*2+1,mid+1,R);
    seg[id] = max(seg[id*2],seg[id*2+1]);
    return;
}

void modify(int id,int L,int R,int pos,int val) {
    if (L==R) {
        seg[id] = val;
        return;
    }
    int mid=(L+R)>>1;
    if (pos <= mid) modify(id*2,L,mid,pos,val);
    else modify(id*2+1,mid+1,R,pos,val);
    seg[id] = max(seg[id*2],seg[id*2+1]);
    return;
}

int query(int id,int L,int R,int l,int r) {
    if (l>R || L>r) return 0;
    else if (l<=L && R<=r) return seg[id];
    int mid=(L+R)>>1;
    return max(query(id*2,L,mid,l,r),query(id*2+1,mid+1,R,l,r));
}

struct Edge {
    int a,b,c;
    void input() {
        scanf("%d %d %d",&a,&b,&c);
    }
} edge[MAX_N];

int sz[MAX_N];
bool visit[MAX_N];
int par[MAX_P][MAX_N];

void dfs1(int id,int p) {
    par[0][id] = p;
    sz[id] = 1;
    visit[id] = 1;
    for (auto i:edg[id]) {
        if (!visit[i]) {
            dfs1(i,id);
            sz[id] += sz[i];
        }
    }
}

int tin[MAX_N],tout[MAX_N];
int ttin[MAX_N];
int head[MAX_N];
int stamp1,stamp2;
int depth[MAX_N];

void dfs2(int id,int headd,int cur_depth) {
    ttin[id] = stamp1;
    tin[id] = stamp2;
    stamp1++;
    stamp2++;
    visit[id]=1;
    head[id] = headd;
    depth[id] = cur_depth;
    sort(edg[id].begin(),edg[id].end(),[](int &a,int &b) {
        return sz[a] > sz[b];
    });
    bool flag=false;
    for (auto i:edg[id]) {
        if (!visit[i]) {
            if (!flag) {
                dfs2(i,headd,cur_depth+1);
                flag=1;
            }
            else {
                dfs2(i,i,cur_depth+1);
            }
        }
    }
    tout[id]=stamp2;
    stamp2++;
}

bool is_anc(int son,int parent) {
    return tin[son]>=tin[parent] && tout[parent]>=tout[son];
}

int get_real_lca(int x,int y) {
    if (depth[x] > depth[y]) swap(x,y);
    if (is_anc(y,x)) return x;
    for (int i=MAX_P-1;i>=0;i--) {
        if (!is_anc(x,par[i][y])) y=par[i][y];
    }
    return par[0][y];
}

int get_fake_lca(int parent,int son) {
    if (parent == son) return 0;
    for (int i=MAX_P-1;i>=0;i--) {
        if (!is_anc(parent,par[i][son])) {
            son = par[i][son];
        }
    }
    return son;
}

int main () {
    int T;
    scanf("%d",&T);
    while (T--) {
        int n;
        scanf("%d",&n);
        for (int i=0;n>=i;i++) {
            edg[i].clear();
        }
        for (int i=1;n>i;i++) {
            edge[i].input();
            edg[edge[i].a].push_back(edge[i].b);
            edg[edge[i].b].push_back(edge[i].a);
        }
        memset(visit,0,sizeof(visit));
        dfs1(1,1);
        memset(visit,0,sizeof(visit));
        stamp1=stamp2=1;
        dfs2(1,1,1);
        for (int i=1;n>i;i++) {
            int &_a=edge[i].a,&_b=edge[i].b,_c=edge[i].c;
            if (ttin[_a] > ttin[_b]) swap(_a,_b);
            val[ttin[_b]] = _c;
        }
        val[1] = 0;
        init(1,1,n);
        for (int i=1;MAX_P>i;i++) {
            for (int j=1;n>=j;j++) {
                par[i][j] = par[i-1][par[i-1][j]];
            }
        }
        char s[10];
        while (1) {
            getchar();
            scanf("%s",s);
            if (s[0]=='D') {
                //huuray!!!
                break;
            }
            else if (s[0]=='C') {
                int a,b;
                scanf("%d %d",&a,&b);
                modify(1,1,n,ttin[edge[a].b],b);
            }
            else {
                int a,b;
                scanf("%d %d",&a,&b);
                if (depth[a] > depth[b]) swap(a,b);
                int lca=get_real_lca(a,b);
                if (a==b) {
                    puts("0");
                    continue;
                }
                int destination=get_fake_lca(lca,b);
                int ans=0;
                while (depth[destination] <= depth[b]) {
                    if (depth[head[b]] >= depth[destination]) {
                        ans = max(ans,query(1,1,n,ttin[head[b]],ttin[b]));
                        b=head[b];
                        b=par[0][b];
                    }
                    else {
                        ans = max(ans,query(1,1,n,ttin[destination],ttin[b]));
                        break;
                    }
                }
                if (lca != a) {
                    destination = get_fake_lca(lca,a);
                    b=a;
                    while (depth[destination] <= depth[b]) {
                        if (depth[head[b]] >= depth[destination]) {
                            ans = max(ans,query(1,1,n,ttin[head[b]],ttin[b]));
                            b=head[b];
                            b=par[0][b];
                        }
                        else {
                            ans = max(ans,query(1,1,n,ttin[destination],ttin[b]));
                            break;
                        }
                    }
                }
                printf("%d\n",ans);
            }
        }
    }
}