SPOJ QTREE - Query on a tree

樹鍊剖分


#include <cstdio>
#include <cstdlib>
#include <algorithm>
#include <cmath>
#include <vector>
#include <cassert>

using namespace std;

struct BST{
    vector<int> tree;
    int base;
    BST(const vector<int>& arr){
        base = (1<<(int)(ceil(log2((int)arr.size() + 4)))) - 1;
        tree = vector<int>(base*2 + 2, 0);
        for(int lx = 0;lx < arr.size();lx++) tree[base+2+lx] = arr[lx];
        for(int lx = base;lx;lx--) tree[lx] = max(tree[lx*2], tree[lx*2+1]);
        return;
    }
   
    void modify(int p, int val){
        p+=base+2; tree[p] = val;
        for(p>>=1;p;p>>=1)
            tree[p] = max(tree[p*2], tree[p*2+1]);
        return;
    }

    int query(int a, int b){
        assert(a <= b);
        int ret = 0;
        for(a+=base+1, b+=base+3;a^b^1;a>>=1, b>>=1){
            if(~a&1) ret = max(ret, tree[a^1]);
            if(b&1) ret = max(ret, tree[b^1]);
        }
        return ret;
    }
};

vector<int> tab[10000];
vector<int> tabwei[10000];

int  siz[10000];
int  dep[10000];
int  fat[10000];
int  nup[10000][15];
int  wei[10000];
BST* bst[10000];
vector<int> lin[10000];

void initial(int n){
    for(int lx = 0;lx < n;lx++) tab[lx].clear(), tabwei[lx].clear();
    return;
}

void finish(int n){
    for(int lx = 0;lx < n;lx++)
        if(fat[lx] == lx)
            delete bst[lx];
    return;
}

void dfs(int v, int ft){
    siz[v] = 1, fat[v] = v;
    int sz = 0, max_nd;
    for(int lx = 0;lx < tab[v].size();lx++){
        int nd = tab[v][lx];
        if(nd == ft) continue;
        nup[nd][0] = v, dep[nd] = dep[v]+1;
        wei[nd] = tabwei[v][lx];
        dfs(nd, v);
        siz[v] += siz[nd];
        if(siz[nd] > sz) sz = siz[nd], max_nd = nd;
    }
    if(sz) fat[max_nd] = v;
    return;
}

int get_fat(int a){return a == fat[a] ? a: fat[a] = get_fat(fat[a]);}

int st_by_dep(int a, int b){return dep[a] < dep[b];}

void make_drec(int n){
    dep[0] = 0;
    dfs(0, 0);

    nup[0][0] = 0;
    for(int lx = 1;lx < 15;lx++)
        for(int ly = 0;ly < n;ly++)
            nup[ly][lx] = nup[nup[ly][lx-1]][lx-1];

    for(int lx = 0;lx < n;lx++) lin[lx].clear();
    for(int lx = 0;lx < n;lx++) lin[get_fat(lx)].push_back(lx);
    for(int lx = 0;lx < n;lx++){
        if(fat[lx] != lx) continue;
        sort(lin[lx].begin(), lin[lx].end(), st_by_dep);
        vector<int> tmp;
        for(int ly = 0; ly < lin[lx].size();ly++)
            tmp.push_back(wei[lin[lx][ly]]);
        bst[lx] = new BST(tmp);
    }
    return;
}

int get_nup(int v, int step){
    for(int lx = 0;lx < 15;lx++)
        if(step&(1<<lx))
            v = nup[v][lx];
    return v;
}

int get_lca(int a, int b){
    if(dep[a] > dep[b]) a = get_nup(a, dep[a]-dep[b]);
    else if(dep[b] > dep[a]) b = get_nup(b, dep[b]-dep[a]);
    int ls = 0, rs = 10000;
    while(ls < rs){
        int ts = (ls+rs)>>1;
        if(get_nup(a, ts) == get_nup(b, ts))
            rs = ts;
        else
            ls = ts+1;
    }
    return get_nup(a, ls);
}

void change_edge(int a, int b, int c){
    // a--> b
    if(dep[a] > dep[b]) swap(a, b);
    bst[fat[b]]->modify(dep[b]-dep[fat[b]], c);
    return;
}

int _query_path(int anc, int chd){
    int ret = 0;
    while(chd != anc){
        int fat_id = fat[chd];
        if(dep[anc] >= dep[fat_id]){
            ret = max(ret, bst[fat_id]->query(dep[anc]-dep[fat_id]+1, dep[chd]-dep[fat_id]));
            break;
        }else{
            ret = max(ret, bst[fat_id]->query(0, dep[chd]-dep[fat_id]));
            chd = nup[fat_id][0];
        }
       
    }
    return ret;
}

int query_path(int a, int b){
    int anc = get_lca(a, b);
    assert(dep[anc] <= dep[a] and dep[anc] <= dep[b]);
    // anc --> a
    // anc --> b
    return max(_query_path(anc, a), _query_path(anc, b));
}

struct edge{int a, b;};

int main(){
    int T; scanf("%d", &T);
    while(T--){
        int n; scanf("%d", &n);
        initial(n);
        vector<edge> es;
        for(int lx = 0;lx < n-1;lx++){
            int a, b, c; scanf("%d %d %d", &a, &b, &c);
            a--, b--;
            tab[a].push_back(b), tabwei[a].push_back(c);
            tab[b].push_back(a), tabwei[b].push_back(c);

            edge _e = {a, b};
            es.push_back(_e);
        }
        make_drec(n);
        for(;;){
            char cmd[10]; scanf("%s", cmd);
            int a, b;
            if(cmd[0] == 'Q'){
                scanf("%d %d", &a, &b);
                a--, b--;
                printf("%d\n", query_path(a, b));
            }else if(cmd[0] == 'C'){
                scanf("%d %d", &a, &b);
                a--;
                change_edge(es[a].a, es[a].b, b);
            }else if(cmd[0] == 'D'){
                break;
            }else{
                assert(0);
            }
        }
        finish(n);
        puts("");
    }
    return 0;
}