对于树剖里的线段树,我们多维护一个当前区间左端点的颜色和当前区间右端点的颜色
所以(pushup)时就变成了这样:
t[p].lc = t[p << 1].lc, t[p].rc = t[p << 1 | 1].rc;
t[p].s = t[p << 1].s + t[p << 1 | 1].s;
if (t[p << 1].rc == t[p << 1 | 1].lc) --t[p].s;
还要注意注意在我们跳链时,假如当前id[top[u]]的颜色与id[fa[top[u]]]的颜色相同,我们要记得给答案减1,因为它们是同种颜色属于同一段而我们把这段切成两段加了两次
#include <cstdio>
#include <algorithm>
const int N = 1e5 + 30;
int n, m, len, cnt, x, a, b, c;
int w[N], f[N], dep[N], fa[N], top[N], siz[N], son[N], dfn[N], id[N];
char ch[5];
struct Edge
{
int to, nx;
}e[N << 1];
struct SegmentTree
{
int l, r, s, lc, rc, lazy;
}t[N << 2];
int read()
{
int s = 0, ff = 1;
char cc = getchar();
while (cc < '0' || cc > '9')
{
if (cc == '-') ff = -1;
cc = getchar();
}
while (cc >= '0' && cc <= '9')
{
s = s * 10 + cc - '0';
cc = getchar();
}
return s * ff;
}
inline void add_edge(int from, int to)
{
e[++len] = (Edge){to, f[from]};
f[from] = len;
}
void dfs1(int p, int from)
{
fa[p] = from, dep[p] = dep[fa[p]] + 1, siz[p] = 1;
for (int i = f[p]; i; i = e[i].nx)
{
if (fa[p] == e[i].to) continue;
dfs1(e[i].to, p), siz[p] += siz[e[i].to];
if (siz[e[i].to] > siz[son[p]]) son[p] = e[i].to;
}
}
void dfs2(int p, int tp)
{
dfn[++cnt] = p, id[p] = cnt, top[p] = tp;
if (son[p]) dfs2(son[p], tp);
for (int i = f[p]; i; i = e[i].nx)
if (e[i].to != son[p] && e[i].to != fa[p]) dfs2(e[i].to, e[i].to);
}
void build(int p, int l, int r)
{
t[p].l = l, t[p].r = r;
if (l == r) { t[p].s = 1; t[p].lc = t[p].rc = w[dfn[l]]; return ; }
int mid = l + r >> 1;
build(p << 1, l, mid), build(p << 1 | 1, mid + 1, r);
t[p].lc = t[p << 1].lc, t[p].rc = t[p << 1 | 1].rc, t[p].s = t[p << 1].s + t[p << 1 | 1].s;
if (t[p << 1].rc == t[p << 1 | 1].lc) --t[p].s;
}
inline void spread(int p)
{
if (!t[p].lazy) return ;
t[p << 1].s = t[p << 1 | 1].s = 1;
t[p << 1].lazy = t[p << 1].lc = t[p << 1].rc = t[p].lazy;
t[p << 1 | 1].lazy = t[p << 1 | 1].lc = t[p << 1 | 1].rc = t[p].lazy;
t[p].lazy = 0;
}
void change(int p, int l, int r, int col)
{
if (t[p].l >= l && t[p].r <= r)
{ t[p].lc = t[p].rc = t[p].lazy = col, t[p].s = 1; return ; }
spread(p);
int mid = t[p].l + t[p].r >> 1;
if (l <= mid) change(p << 1, l, r, col);
if (mid < r) change(p << 1 | 1, l, r, col);
t[p].lc = t[p << 1].lc, t[p].rc = t[p << 1 | 1].rc, t[p].s = t[p << 1].s + t[p << 1 | 1].s;
if (t[p << 1].rc == t[p << 1 | 1].lc) --t[p].s;
}
int ask(int p, int l, int r)
{
if (t[p].l >= l && t[p].r <= r) return t[p].s;
spread(p);
int mid = t[p].l + t[p].r >> 1, tot = 0;
if (l <= mid && mid < r && t[p << 1].rc == t[p << 1 | 1].lc) --tot;
if (l <= mid) tot += ask(p << 1, l, r);
if (mid < r) tot += ask(p << 1 | 1, l, r);
return tot;
}
int get(int p, int pos)
{
if (t[p].l == t[p].r) return t[p].lc;
spread(p);
int mid = t[p].l + t[p].r >> 1;
if (pos <= mid) return get(p << 1, pos);
else return get(p << 1 | 1, pos);
}
inline void treechange(int u, int v, int col)
{
while (top[u] != top[v])
{
if (dep[top[u]] < dep[top[v]]) std::swap(u, v);
change(1, id[top[u]], id[u], col);
u = fa[top[u]];
}
if (dep[u] < dep[v]) std::swap(u, v);
change(1, id[v], id[u], col);
}
inline int treeask(int u, int v)
{
int ans = 0, cu, cv;
while (top[u] != top[v])
{
if (dep[top[u]] < dep[top[v]]) std::swap(u, v);
ans += ask(1, id[top[u]], id[u]);
cu = get(1, id[top[u]]), cv = get(1, id[fa[top[u]]]);
if (cu == cv) --ans;
u = fa[top[u]];
}
if (dep[u] < dep[v]) std::swap(u, v);
ans += ask(1, id[v], id[u]);
if (ans == 0) return 1;
return ans;
}
int main()
{
n = read(), m = read();
for (int i = 1; i <= n; ++i) w[i] = read();
for (int i = 1; i < n; ++i)
{
a = read(), b = read();
add_edge(a, b), add_edge(b, a);
}
dfs1(1, 0), dfs2(1, 1), build(1, 1, n);
for (int i = 1; i <= m; ++i)
{
scanf ("%s", ch);
if (ch[0] == 'C')
{
a = read(), b = read(), c = read();
treechange(a, b, c);
}
else
{
a = read(), b = read();
printf ("%d
", treeask(a, b));
}
}
return 0;
}