HYSBZ 2243 染色 （树链剖分）

HYSBZ 2243 染色

#include <cstdio>
#include <cstring>
#include <vector>
#include <algorithm>
using namespace std;

const int N = 100005;

int dep
, fa
, son
, sz
, top
, id
, idx, val
, val2
;
int first
, next[N * 2], vv[N * 2], en;

void init() {
en = 0;
idx = 0;
memset(first, -1, sizeof(first));
}

void add_Edge(int u, int v) {
vv[en] = v;
next[en] = first[u];
first[u] = en++;
}

void dfs1(int u, int f, int d) {
dep[u] = d;
sz[u] = 1;
fa[u] = f;
son[u] = 0;
for (int i = first[u]; i + 1; i = next[i]) {
int v = vv[i];
if (v == f) continue;
dfs1(v, u, d + 1);
sz[u] += sz[v];
if (sz[son[u]] < sz[v])
son[u] = v;
}
}

void dfs2(int u, int tp) {
id[u] = ++idx;
top[u] = tp;
if (son[u]) dfs2(son[u], tp);
for (int i = first[u]; i + 1; i = next[i]) {
int v = vv[i];
if (v == fa[u] || v == son[u]) continue;
dfs2(v, v);
}
}

#define lson(x) ((x<<1)+1)
#define rson(x) ((x<<1)+2)

struct Node {
int l, r, lc, rc, cnt, setv;
Node() {cnt = 0;}
} node[N * 4];

void pushup(int x) {
node[x].cnt = node[lson(x)].cnt + node[rson(x)].cnt;
if (node[lson(x)].rc == node[rson(x)].lc) node[x].cnt--;
node[x].lc = node[lson(x)].lc;
node[x].rc = node[rson(x)].rc;
}

void pushdown(int x) {
if (node[x].setv == -1) return;
node[lson(x)].setv = node[lson(x)].lc = node[lson(x)].rc = node[x].setv;
node[rson(x)].setv = node[rson(x)].lc = node[rson(x)].rc = node[x].setv;
node[lson(x)].cnt = node[rson(x)].cnt = 1;
node[x].setv = -1;
}

void build(int l, int r, int x = 0) {
node[x].l = l; node[x].r = r;
node[x].setv = -1;
if (l == r) {
node[x].lc = node[x].rc = val[l];
node[x].cnt = 1;
return;
}
int mid = (l + r) / 2;
build(l, mid, lson(x));
build(mid + 1, r, rson(x));
pushup(x);
}

void add(int l, int r, int v, int x = 0) {
if (node[x].l >= l && node[x].r <= r) {
node[x].cnt = 1;
node[x].setv = node[x].lc = node[x].rc = v;
return;
}
pushdown(x);
int mid = (node[x].l + node[x].r) / 2;
if (l <= mid) add(l, r, v, lson(x));
if (r > mid) add(l, r, v, rson(x));
pushup(x);
}

Node merge(Node a, Node b) {
Node ans;
if (a.cnt == 0) return b;
if (b.cnt == 0) return a;
ans.lc = a.lc;
ans.rc = b.rc;
ans.cnt = a.cnt + b.cnt;
if (a.rc == b.lc) ans.cnt--;
return ans;
}

Node query(int l, int r, int x = 0) {
if (node[x].l >= l && node[x].r <= r) {
return node[x];
}
pushdown(x);
int mid = (node[x].l + node[x].r) / 2;
Node ans;
if (l <= mid) ans = merge(query(l, r, lson(x)), ans);
if (r > mid) ans = merge(ans, query(l, r, rson(x)));
pushup(x);
return ans;
}

void gao1(int u, int v, int val) {
int tp1 = top[u], tp2 = top[v];
while (tp1 != tp2) {
if (dep[tp1] < dep[tp2]) {
swap(tp1, tp2);
swap(u, v);
}
add(id[tp1], id[u], val);
u = fa[tp1];
tp1 = top[u];
}
if (dep[u] > dep[v]) swap(u, v);
add(id[u], id[v], val);
}

int gao2(int u, int v) {
int tp1 = top[u], tp2 = top[v];
Node ans1, ans2;
while (tp1 != tp2) {
if (dep[tp1] < dep[tp2]) {
swap(tp1, tp2);
swap(u, v);
swap(ans1, ans2);
}
ans1 = merge(query(id[tp1], id[u]), ans1);
u = fa[tp1];
tp1 = top[u];
}
if (dep[u] > dep[v]) {
swap(ans1, ans2);
swap(u, v);
}
ans2 = merge(query(id[u], id[v]), ans2);
if (ans1.cnt == 0) return ans2.cnt;
if (ans2.cnt == 0) return ans1.cnt;
int ans = ans1.cnt + ans2.cnt;
if (ans1.lc == ans2.lc) ans--;
return ans;
}

int n, m;

int main() {
while (~scanf("%d%d", &n, &m)) {
init();
for (int i = 1; i <= n; i++) scanf("%d", &val2[i]);
int u, v;
for (int i = 1; i < n; i++) {
scanf("%d%d", &u, &v);
add_Edge(u, v);
add_Edge(v, u);
}
dfs1(1, 0, 1);
dfs2(1, 1);
for (int i = 1; i <= n; i++)
val[id[i]] = val2[i];
build(1, idx);
char Q[10];
int a, b, c;
while (m--) {
scanf("%s%d%d", Q, &a, &b);
if (Q[0] == 'Q') printf("%d\n", gao2(a, b));
else {
scanf("%d", &c);
gao1(a, b, c);
}
}
}
return 0;
}