题意:中文题。
思路:进行树链剖分后用线段树维护区间,注意区间合并时相邻端点颜色相同时ans-1,我刚开始对区间端点的处理是建立在原数组上的,改成结构体储存区间左右端点颜色就A了。
代码:
#include <iostream>
using namespace std;
typedef long long ll;
const int N = 3e5 + 10;
struct edge
{
int to, next;
}Edge[N<<1];
struct node
{
int sum, lc, rc;
}Node[N<<2];
int n, m, a[N], head[N], lazy[N<<2];
int cnt, f[N], dep[N], size[N], son[N], rk[N], top[N], id[N];
//树链剖分模板
void add_edge(int A,int B)
{
Edge[++cnt].to = B;
Edge[cnt].next = head[A];
head[A] = cnt;
}
void dfs1(int v,int fa,int depth)
{
f[v] = fa;
dep[v] = depth;
size[v] = 1;
for(int i = head[v]; i; i = Edge[i].next){
int to = Edge[i].to;
if(to == fa) continue;
dfs1(to,v,depth+1);
size[v] += size[to];
if(size[to] > size[son[v]])
son[v] = to;
}
}
void dfs2(int v,int tp)
{
top[v] = tp;
id[v] = ++cnt;
rk[cnt] = v;
if(!son[v]) return ;
dfs2(son[v],tp);
for(int i = head[v]; i; i = Edge[i].next){
int to = Edge[i].to;
if(to != son[v] && to != f[v])
dfs2(to,to);
}
}
void pushUp(int rt)
{
Node[rt].sum = Node[rt<<1].sum + Node[rt<<1|1].sum - (Node[rt<<1].rc == Node[rt<<1|1].lc);
Node[rt].lc = Node[rt<<1].lc, Node[rt].rc = Node[rt<<1|1].rc;
}
void pushDown(int rt)
{
if(lazy[rt] == -1) return ;
Node[rt<<1].sum = Node[rt<<1|1].sum = 1;
Node[rt<<1].lc = Node[rt<<1].rc = Node[rt<<1|1].lc = Node[rt<<1|1].rc = lazy[rt];
lazy[rt<<1] = lazy[rt<<1|1] = lazy[rt];
lazy[rt] = -1;
}
void build(int l,int r,int rt)
{
lazy[rt] = -1;
if(l == r){
Node[rt].sum = 1;
Node[rt].lc = Node[rt].rc = a[rk[l]];
return ;
}
int m = (l+r)>>1;
build(l,m,rt<<1);
build(m+1,r,rt<<1|1);
pushUp(rt);
}
void update(int L,int R,int c,int l,int r,int rt)
{
if(L <= l && r <= R){
Node[rt].sum = 1;
Node[rt].lc = Node[rt].rc = c;
lazy[rt] = c;
return ;
}
int m = (l+r)>>1;
pushDown(rt);
if(L <= m) update(L,R,c,l,m,rt<<1);
if(m < R) update(L,R,c,m+1,r,rt<<1|1);
pushUp(rt);
}
int query(int L,int R,int l,int r,int rt)
{
if(L <= l && r <= R){
return Node[rt].sum;
}
int m = (l+r)>>1;
ll ans = 0, flag = 0;
pushDown(rt);
if(L <= m) ans += query(L,R,l,m,rt<<1), flag += 1;
if(m < R) ans += query(L,R,m+1,r,rt<<1|1), flag += 1;
if(flag == 2) ans -= (Node[rt<<1].rc == Node[rt<<1|1].lc);
return ans;
}
int getcolor(int L,int l,int r,int rt)
{
if(l == r) return Node[rt].lc;
pushDown(rt);
int m = (l+r)>>1;
if(L <= m) return getcolor(L,l,m,rt<<1);
else return getcolor(L,m+1,r,rt<<1|1);
}
void updates(int A,int B,int C)
{
while(top[A] != top[B]){
if(dep[top[A]] < dep[top[B]]) swap(A,B);
update(id[top[A]],id[A],C,1,n,1);
A = f[top[A]];
}
if(id[A] > id[B]) swap(A,B);
update(id[A],id[B],C,1,n,1);
}
ll querys(int A,int B)
{
ll ans = 0;
while(top[A] != top[B]){
if(dep[top[A]] < dep[top[B]]) swap(A,B);
ans += query(id[top[A]],id[A],1,n,1);
if(getcolor(id[top[A]],1,n,1) == getcolor(id[f[top[A]]],1,n,1)) ans -= 1;
A = f[top[A]]; //在这里对所有lazy标记进行下移,如果该链的头节点到它的父节点的颜色相同时,答案减一
}
if(id[A] > id[B]) swap(A,B);
return ans + query(id[A],id[B],1,n,1);
}
int main()
{
ios::sync_with_stdio(false);
cin.tie(0), cout.tie(0);
int A, B, C;
char oper;
cin>>n>>m;
for(int i = 1; i <= n; ++i)
cin>>a[i];
for(int i = 0; i < n - 1; ++i){
cin>>A>>B;
add_edge(A,B); add_edge(B,A);
}
cnt = 0;
dfs1(1,0,1);
dfs2(1,1);
build(1,n,1);
while(m--){
cin>>oper;
if(oper == 'Q'){
cin>>A>>B;
cout<<querys(A,B)<<'\n';
}
else{
cin>>A>>B>>C;
updates(A,B,C);
}
}
return 0;
}