LightOJ - 1348 Aladdin and the Return Journey(树剖)

题目大意：给出一棵树和两种操作
0 i j：从i点到j点的权值和
1 i v: 第i点的权值变成v

解题思路：树剖裸题

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
const int N = 30010;
const int M = 30010 << 2;

struct Edge{
    int u, v, next;
    Edge() {}
    Edge(int u, int v, int next): u(u), v(v), next(next) {}
}E[N * 2];

int size[N], son[N], top[N], id[N], head[N], val[N], id2[N], dep[N], f[N];
int left[M], right[M], sum[M];
int idx, n, tot, cas = 1;

void build(int u, int l, int r) {
    left[u] = l; right[u] = r;
    if (l == r) {
        sum[u] = val[id2[l]];
        return ;
    }
    int mid = (l + r) >> 1;
    build(2 * u, l, mid);
    build(2 * u + 1, mid + 1, r);
    sum[u] = sum[2 * u] + sum[2 * u + 1];
}


void dfs1(int u, int fa, int depth) {
    dep[u] = depth; f[u] = fa; size[u] = 1; son[u] = 0;
    for (int i = head[u]; ~i; i = E[i].next) {
        int v = E[i].v;
        if (v == fa) continue;
        dfs1(v, u, depth + 1);
        size[u] += size[v];
        if (size[son[u]] < size[v]) son[u] = v;
    }
}

void dfs2(int u, int tp) {
    top[u] = tp;
    id[u] = ++idx;
    if (son[u]) dfs2(son[u], tp);
    for (int i = head[u]; ~i; i = E[i].next) {
        int v = E[i].v;
        if (v == f[u] || v == son[u]) continue;
        dfs2(v, v);
    }
}

void AddEdge(int u, int v) {
    E[tot] = Edge(u, v, head[u]);
    head[u] = tot++;
    E[tot] = Edge(v, u, head[v]);
    head[v] = tot++;
}

int query(int u, int l, int r) {
    if (l <= left[u] && right[u] <= r) return sum[u];

    int mid = (left[u] + right[u]) >> 1;
    int ans = 0;
    if (l <= mid) ans += query(2 * u, l, r);
    if (r > mid) ans +=  query(2 * u + 1, l, r);

    return ans;
}

int gao(int u, int v) {
    int tp1 = top[u], tp2 = top[v];
    int ans = 0;
    while (tp1 != tp2) {
        if (dep[tp1] < dep[tp2]) {
            swap(tp1, tp2);
            swap(u, v);
        }

        ans += query(1, id[tp1], id[u]);
        u = f[tp1];
        tp1 = top[u];
    }
    if (dep[u] > dep[v]) swap(u, v);
    ans += query(1, id[u], id[v]);
    return ans;
}

void modify(int u, int pos, int c) {
    if (left[u] == right[u] && left[u] == pos) {
        sum[u] = c;
        return ;
    }
    int mid = (left[u] + right[u]) >> 1;
    if (pos <= mid) modify(2 * u, pos, c);
    else modify(2 * u + 1, pos, c); 

    sum[u] = sum[2 * u + 1] + sum[2 * u];
}

void solve() {
    printf("Case %d:\n", cas++);
    int q, op, a, b;
    scanf("%d", &q);
    while (q--) {
        scanf("%d%d%d", &op, &a, &b);
        if (op) modify(1, id[a + 1], b);
        else printf("%d\n", gao(a + 1, b + 1));
    }
}

void init() {
    scanf("%d", &n);
    memset(head, -1, sizeof(head));
    tot = 0;

    for (int i = 1; i <= n; i++)
        scanf("%d", &val[i]);
    int u, v;
    for (int i = 1; i < n; i++) {
        scanf("%d%d", &u, &v);
        AddEdge(u + 1, v + 1);
    }


    idx = 0;
    dfs1(1, 0, 1);
    dfs2(1, 1);
    for (int i = 1; i <= n; i++) id2[id[i]] = i;
    build(1, 1, idx);
}


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