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

#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
, son
, top
, id
, head
, val
, id2
, dep
, f
;
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;
}