树链剖分+线段树 BZOJ 1036 [ZJOI2008]树的统计Count

#include <bits/stdc++.h>

const int N = 3e4 + 5;
const int D = 20;
const int INF = 0x7fffffff;
struct Edge {
int v, nex;
}edge[N<<1];
int head
;
int a
;
int sz
, dep
, belong
, pos
;
int rt
[D];
int n, tote, loc;

void add_edge(int u, int v) {
edge[tote].v = v; edge[tote].nex = head[u];
head[u] = tote++;
}

void init_edge() {
memset (head, -1, sizeof (head));
tote = 0;
}

void DFS1(int u, int fa) {
sz[u] = 1; dep[u] = dep[fa] + 1;
rt[u][0] = fa;
for (int i=head[u]; ~i; i=edge[i].nex) {
int v = edge[i].v;
if (v == fa) {
continue;
}
DFS1 (v, u);
sz[u] += sz[v];
}
}

void DFS2(int u, int fa, int chain) {
int k = 0;
pos[u] = ++loc;
belong[u] = chain;
for (int i=head[u]; ~i; i=edge[i].nex) {
int v = edge[i].v;
if (v == fa) {
continue;
}
if (sz[v] > sz[k]) {
k = v;
}
}
if (k == 0) {
return ;
}
DFS2 (k, u, chain);
for (int i=head[u]; ~i; i=edge[i].nex) {
int v = edge[i].v;
if (v == fa || v == k) {
continue;
}
DFS2 (v, u, v);
}
}

void init_LCA() {
for (int j=1; j<D; ++j) {
for (int i=1; i<=n; ++i) {
rt[i][j] = rt[i][j-1] == 0 ? 0 : rt[rt[i][j-1]][j-1];
}
}
}

int LCA(int u, int v) {
if (dep[u] < dep[v]) {
std::swap (u, v);
}
for (int i=0; i<D; ++i) {
if ((dep[u] - dep[v]) >> i & 1) {
u = rt[u][i];
}
}
if (u == v) {
return u;
}
for (int i=D-1; i>=0; --i) {
if (rt[u][i] != rt[v][i]) {
u = rt[u][i];
v = rt[v][i];
}
}
return rt[u][0];
}

int mx[N<<2], sum[N<<2];

#define lson l, mid, o << 1
#define rson mid + 1, r, o << 1 | 1

void push_up(int o) {
mx[o] = std::max (mx[o<<1], mx[o<<1|1]);
sum[o] = sum[o<<1] + sum[o<<1|1];
}

void updata(int p, int v, int l, int r, int o) {
if (l == r) {
mx[o] = sum[o] = v;
return ;
}
int mid = l + r >> 1;
if (p <= mid) {
updata (p, v, lson);
} else {
updata (p, v, rson);
}
push_up (o);
}

int query_max(int ql, int qr, int l, int r, int o) {
if (ql <= l && r <= qr) {
return mx[o];
}
int mid = l + r >> 1, ret = -INF;
if (ql <= mid) {
ret = std::max (ret, query_max (ql, qr, lson));
}
if (qr > mid) {
ret = std::max (ret, query_max (ql, qr, rson));
}
return ret;
}

int query_sum(int ql, int qr, int l, int r, int o) {
if (ql <= l && r <= qr) {
return sum[o];
}
int mid = l + r >> 1, ret = 0;
if (ql <= mid) {
ret += query_sum (ql, qr, lson);
}
if (qr > mid) {
ret += query_sum (ql, qr, rson);
}
return ret;
}

int get_sum(int u, int lca) {
int ret = 0;
while (belong[u] != belong[lca]) {
ret += query_sum (pos[belong[u]], pos[u], 1, n, 1);
u = rt[belong[u]][0];
}
ret += query_sum (pos[lca], pos[u], 1, n, 1);
return ret;
}

int get_max(int u, int lca) {
int ret = -INF;
while (belong[u] != belong[lca]) {
ret = std::max (ret, query_max (pos[belong[u]], pos[u], 1, n, 1));
u = rt[belong[u]][0];
}
ret = std::max (ret, query_max (pos[lca], pos[u], 1, n, 1));
return ret;
}

void prepare() {
sz[0] = 0; dep[0] = 0;
DFS1 (1, 0);
loc = 0;
DFS2 (1, 0, 1);
init_LCA ();
for (int i=1; i<=n; ++i) {
updata (pos[i], a[i], 1, n, 1);
}
}

int main() {
scanf ("%d", &n);
init_edge ();
for (int i=1; i<n; ++i) {
int u, v;
scanf ("%d%d", &u, &v);
add_edge (u, v);
add_edge (v, u);
}
for (int i=1; i<=n; ++i) {
scanf ("%d", a+i);
}

prepare ();

int q;
scanf ("%d", &q);
char ch[6];
int x, y;
while (q--) {
scanf ("%s%d%d", ch, &x, &y);
if (ch[0] == 'C') {
a[x] = y;
updata (pos[x], a[x], 1, n, 1);
} else {
int lca = LCA (x, y);
if (ch[1] == 'M') {
printf ("%d\n", std::max (get_max (x, lca), get_max (y, lca)));
} else {
printf ("%d\n", get_sum (x, lca) + get_sum (y, lca) - a[lca]);
}
}
}
return 0;
}