LightOJ 1348 （树链剖分 + 线段树（树状数组））

题目

分析

Code

#include <bits/stdc++.h>

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;
}