LightOJ 1348(트리 체인 분할 + 세그먼트 트리(트리 배열)

#include 

using namespace std;

const int maxn = 30000 + 131;
struct Edge {
    int Next;
    int To;
}edge[maxn<<1];
int Head[maxn], tot, n; 
///         
int top[maxn];      //     
int deep[maxn];     //       
int Pre[maxn];      //   
int size[maxn];     //      
int son[maxn];      //         
///              
int t_s[maxn];      //          
int s_t[maxn];      //          。
int pos;
//    
int w[maxn];


void INIT() {
    tot = pos = 0;
    memset(son, -1, sizeof(son));
    memset(Head,-1, sizeof(Head));
}
////
void Addedge(int from, int to) {
    edge[tot].To   = to;
    edge[tot].Next = Head[from];
    Head[from]     = tot++;
}

void Getlist(int root, int pre, int d) { //    
    deep[root] = d;
    Pre[root]  = pre;
    size[root] = 1;
    for(int i = Head[root]; ~i; i = edge[i].Next) {
        int v = edge[i].To;
        if(v != pre) {
            Getlist(v, root, d+1);
            size[root] += size[v]; //  size
            if(son[root] == -1 || size[son[root]] < size[v])
                son[root] = v; //       
        }
    }
}
////        
void Lisan_TtoS(int u, int root) {
    top[u]      = root;
    t_s[u]      = ++pos;
    s_t[t_s[u]] = u;
    if(son[u] == -1) return ;
    Lisan_TtoS(son[u], root);

    for(int i = Head[u]; ~i; i = edge[i].Next) {
        int v = edge[i].To;
        if(v != son[u] && v != Pre[u])
            Lisan_TtoS(v,v); //      .
    }
}
////   
int Sum[maxn << 2];
#define lson l, m, rt<<1
#define rson m+1, r, rt<<1|1

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] = w[s_t[l]];
        /*cout << "This is tree idx:" << s_t[l] \
             << " values:" << w[s_t[l]] << endl;*/
        return ;
    }
    int m = (l + r) >> 1;
    Build(lson);
    Build(rson);
    PushUp(rt);
}

void Update(int pos, int val, int l, int r, int rt) {
    if(l == r) {
        Sum[rt] = val;
        //cout << "This is the tree id: " << s_t[l] << endl;
        return ;
    }
    int m = (l + r) >> 1;
    if(pos <= m) Update(pos, val, lson);
    else Update(pos, val, rson);
    PushUp(rt);
}

int Query(int L, int R, int l, int r, int rt) {
    //cout << "seg l : " << l << "  r : " << r ;
    //cout << "   Find L:" << L << "  R: " << R << endl;
    if(L <= l && r <= R) {
        //cout << "Had add : " << Sum[rt] << endl;
        return Sum[rt];
    }
    int m = (l + r) >> 1;
    int ret = 0;
    if(L <= m) ret += Query(L, R, lson);
    if(R >  m) ret += Query(L, R, rson);
    return ret;
}
////  （u，v）
int Find(int u, int v) {    /// u to v
    int fa_u = top[u];      /// u      .
    int fa_v = top[v];
    int ret = 0;
    while(fa_u != fa_v) {
        if(deep[fa_u] < deep[fa_v]) {
            swap(u, v);
            swap(fa_v,fa_u);
        }
        //cout << "This is the list :" << fa_u << "->" << u << endl; 
        //cout << "This is the Segm :" << t_s[fa_u] << "->" << t_s[u] << endl;
        ret += Query(t_s[fa_u], t_s[u], 1, n, 1);
        u = Pre[fa_u];
        fa_u = top[u];
    }
    //  ，            。
    //     ，      i   to j     v    j     
    // root          or (0) or      。
    if(deep[u] > deep[v]) swap(u, v);
    ret += Query(t_s[u], t_s[v], 1, n, 1);
    return ret;
}

int main() {
    int T;
    scanf("%d",&T);
    for(int kase = 1; kase <= T; ++kase) {
        scanf("%d",&n);
        INIT();
        for(int i = 1; i <= n; ++i)
            scanf("%d",w+i);
        int u, v;
        for(int i = 1; i < n; ++i) {
            scanf("%d%d",&u, &v);
            u++, v++;
            Addedge(u, v);
            Addedge(v, u);
        }
        Getlist(1, -1, 0);
        Lisan_TtoS(1,1);
        Build(1, n, 1);
        printf("Case %d:
",kase);
        int q;
        scanf("%d",&q);
        for(int i = 0; i < q; ++i) {
            int op;
            scanf("%d%d%d",&op,&u,&v);
            if(op == 1) {
                Update(t_s[++u], v, 1, n, 1);
            }
            else {
                u++, v++;
                printf("%d
", Find(u, v));
            }
        }
    }
    return 0;
}