题意:有一棵树,n个节点,n-1条边,每个点都有一个权值,现在有两种操作,Q a b,询问a到b路径上所有点组成的序列一共有几段,比如11223就是三段,C a b c把节点a到b路径上所有经过的点权值设为c。
题解:线段树的区间合并问题,用线段树维护每个区间的左端点值和右端点值还有区间内有几段。
#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
const int N = 100005;
struct Edge {
int u, v, nxt;
Edge() {}
Edge(int a, int b, int c): u(a), v(b), nxt(c) {}
}e[N << 1];
int val[N], val2[N], size[N], son[N], top[N], id[N], dep[N], fa[N];
int n, m, cnt, tot, head[N], tree[N << 2], flag[N << 2], lv[N << 2], rv[N << 2];
void AddEdge(int u, int v) {
e[cnt] = Edge(u, v, head[u]);
head[u] = cnt++;
e[cnt] = Edge(v, u, head[v]);
head[v] = cnt++;
}
void dfs2(int u, int tp) {
top[u] = tp;
id[u] = ++tot;
if (son[u]) dfs2(son[u], tp);
for (int i = head[u]; i + 1; i = e[i].nxt) {
int v = e[i].v;
if (v == fa[u] || v == son[u]) continue;
dfs2(v, v);
}
}
void dfs1(int u, int f, int depth) {
size[u] = 1, son[u] = 0, dep[u] = depth, fa[u] = f;
for (int i = head[u]; i + 1; i = e[i].nxt) {
int v = e[i].v;
if (v == f) continue;
dfs1(v, u, depth + 1);
size[u] += size[v];
if (size[v] > size[son[u]]) son[u] = v;
}
}
void pushup(int k) {
lv[k] = lv[k * 2];
rv[k] = rv[k * 2 + 1];
tree[k] = tree[k * 2] + tree[k * 2 + 1];
if (rv[k * 2] == lv[k * 2 + 1])
tree[k]--;
}
void pushdown(int k) {
if (flag[k] >= 0) {
tree[k * 2] = tree[k * 2 + 1] = 1;
lv[k * 2] = lv[k * 2 + 1] = flag[k];
rv[k * 2] = rv[k * 2 + 1] = flag[k];
flag[k * 2] = flag[k * 2 + 1] = flag[k];
flag[k] = -1;
}
}
void build(int k, int left, int right) {
flag[k] = -1;
if (left == right) {
lv[k] = rv[k] = val2[left];
tree[k] = 1;
return;
}
int mid = (left + right) / 2;
build(k * 2, left, mid);
build(k * 2 + 1, mid + 1, right);
pushup(k);
}
void modify(int k, int left, int right, int l, int r, int v) {
if (l <= left && right <= r) {
lv[k] = rv[k] = flag[k] = v;
tree[k] = 1;
return;
}
pushdown(k);
int mid = (left + right) / 2;
if (l <= mid)
modify(k * 2, left, mid, l, r, v);
if (r > mid)
modify(k * 2 + 1, mid + 1, right, l, r, v);
pushup(k);
}
int query(int k, int left, int right, int l, int r) {
if (l <= left && right <= r)
return tree[k];
pushdown(k);
int mid = (left + right) / 2;
if (r <= mid)
return query(k * 2, left, mid, l, r);
if (l > mid)
return query(k * 2 + 1, mid + 1, right, l, r);
int res = query(k * 2, left, mid, l, mid) + query(k * 2 + 1, mid + 1, right, mid + 1, r);
if (rv[k * 2] == lv[k * 2 + 1])
res--;
return res;
}
int query2(int k, int left, int right, int pos) {
if (left == right)
return lv[k];
pushdown(k);
int mid = (left + right) / 2;
if (pos <= mid)
return query2(k * 2, left, mid, pos);
return query2(k * 2 + 1, mid + 1, right, pos);
}
void solve(int u, int v, int w, int flag1) {
int tp1 = top[u], tp2 = top[v];
int res = 0;
while (tp1 != tp2) {
if (dep[tp1] < dep[tp2]) {
swap(tp1, tp2);
swap(u, v);
}
if (!flag1) {
res += query(1, 1, tot, id[tp1], id[u]);
int temp1 = query2(1, 1, tot, id[fa[tp1]]);
int temp2 = query2(1, 1, tot, id[tp1]);
if (temp1 == temp2)
res--;
}
else modify(1, 1, tot, id[tp1], id[u], w);
u = fa[tp1];
tp1 = top[u];
}
if (dep[u] > dep[v]) swap(u, v);
if (!flag1)
res += query(1, 1, tot, id[u], id[v]);
else modify(1, 1, tot, id[u], id[v], w);
if (!flag1) printf("%d\n", res);
}
int main() {
while (scanf("%d%d", &n, &m) == 2) {
memset(head, -1, sizeof(head));
cnt = tot = 0;
for (int i = 1; i <= n; i++)
scanf("%d", &val[i]);
int u, v;
for (int i = 0; i < n - 1; i++) {
scanf("%d%d", &u, &v);
AddEdge(u, v);
}
dfs1(1, 0, 1);
dfs2(1, 1);
for (int i = 1; i <= n; i++)
val2[id[i]] = val[i];
build(1, 1, tot);
char op[5];
int a, b, c;
while (m--) {
scanf("%s%d%d", op, &a, &b);
if (op[0] == 'Q') solve(a, b, 0, 0);
else {
scanf("%d", &c);
solve(a, b, c, 1);
}
}
}
return 0;
}