Query on a tree

問題文

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

解法

重軽分解して、鎖をセグ木で管理する。

コード

#include<cstdio>
#include<algorithm>
#include<vector>
#include<cstring>
using namespace std;

//#define int long long

typedef pair<int,int>pint;
typedef vector<int>vint;
typedef vector<pint>vpint;
#define pb push_back
#define mp make_pair
#define fi first
#define se second
#define all(v) (v).begin(),(v).end()
#define rep(i,n) for(int i=0;i<(n);i++)
#define reps(i,f,n) for(int i=(f);i<(n);i++)
#define each(it,v) for(__typeof((v).begin()) it=(v).begin();it!=(v).end();it++)
template<class T,class U>void chmin(T &t,U f){if(t>f)t=f;}
template<class T,class U>void chmax(T &t,U f){if(t<f)t=f;}

const int SIZE=10000;

struct segtree{
    static const int SEG=1<<17;
    vint ma;
    
    void init(){
        ma.assign(SEG*2,0);
    }
    
    void update(int k,int x){
        k+=SEG-1;
        ma[k]=x;
        while(k){
            k=(k-1)/2;
            ma[k]=max(ma[k*2+1],ma[k*2+2]);
        }
    }
    int get(int a,int b,int k=0,int l=0,int r=SEG){
        if(r<=a||b<=l)return 0;
        if(a<=l&&r<=b)return ma[k];
        return max(get(a,b,k*2+1,l,(l+r)/2),get(a,b,k*2+2,(l+r)/2,r));
    }
};

int N;
int A[SIZE],B[SIZE],C[SIZE];
vint G[SIZE];

int par[SIZE],dep[SIZE],hev[SIZE],pos[SIZE],head[SIZE];
segtree seg;

int dfs(int v,int p,int d){
    par[v]=p;
    dep[v]=d;
    hev[v]=-1;
    int ma=0,sum=1;
    rep(i,G[v].size()){
        int to=G[v][i];
        if(to==p)continue;
        int tmp=dfs(to,v,d+1);
        if(ma<tmp){
            ma=tmp;
            hev[v]=to;
        }
        sum+=tmp;
    }
    return sum;
}


void init(){
    dfs(0,-1,0);
    int idx=0;
    rep(i,N){
        if(~par[i]&&hev[par[i]]==i)continue;
        for(int j=i;~j;j=hev[j]){
            head[j]=i;
            pos[j]=idx++;
        }
    }
    seg.init();
    rep(i,N-1){
        if(dep[A[i]]<dep[B[i]])swap(A[i],B[i]);
        seg.update(pos[A[i]],C[i]);
    }
}

int getquery(int a,int b){
    int ret=0;
    while(head[a]!=head[b]){
        if(dep[head[a]]<dep[head[b]])swap(a,b);
        chmax(ret,seg.get(pos[head[a]],pos[a]+1));
        a=par[head[a]];
    }
    if(dep[a]>dep[b])swap(a,b);
    return max(ret,seg.get(pos[a]+1,pos[b]+1));
}

void solve(){
    scanf("%d",&N);
    rep(i,N)G[i].clear();

    rep(i,N-1){
        scanf("%d%d%d",&A[i],&B[i],&C[i]);
        A[i]--;B[i]--;
        G[A[i]].pb(B[i]);
        G[B[i]].pb(A[i]);
    }

    init();

    char type[11];
    while(scanf("%s",type),type[0]!='D'){
        if(type[0]=='Q'){
            int a,b;
            scanf("%d%d",&a,&b);
            a--;b--;
            printf("%d\n",getquery(a,b));
        }
        else{
            int k,x;
            scanf("%d%d",&k,&x);k--;
            seg.update(pos[A[k]],x);
        }
    }
}

signed main(){
    int T;
    scanf("%d",&T);
    while(T--)solve();
    return 0;
}