HYSBZ 2243 树链剖分（区间更新，区间查询）较难

http://www.lydsy.com/JudgeOnline/problem.php?id=2243

Description

给定一棵有n个节点的无根树和m个操作，操作有2类：

1、将节点a到节点b路径上所有点都染成颜色c；

2、询问节点a到节点b路径上的颜色段数量（连续相同颜色被认为是同一段），如“112221”由3段组成：“11”、“222”和“1”。

请你写一个程序依次完成这m个操作。

Input

第一行包含2个整数n和m，分别表示节点数和操作数；

第二行包含n个正整数表示n个节点的初始颜色

下面 行每行包含两个整数x和y，表示x和y之间有一条无向边。

下面 行每行描述一个操作：

“C a b c”表示这是一个染色操作，把节点a到节点b路径上所有点（包括a和b）都染成颜色c；

“Q a b”表示这是一个询问操作，询问节点a到节点b（包括a和b）路径上的颜色段数量。

Output

对于每个询问操作，输出一行答案。

Sample Input

6 5

2 2 1 2 1 1

1 2

1 3

2 4

2 5

2 6

Q 3 5

C 2 1 1

Q 3 5

C 5 1 2

Q 3 5

Sample Output

3

1

2

HINT

数N<=10^5，操作数M<=10^5，所有的颜色C为整数且在[0, 10^9]之间。

/***
HYSBZ 2243 树链剖分（区间更新，区间查询）
解题思路：这道题转化成线段树后需要维护三个值：区间左边界的颜色，区间右边界的颜色，区间的颜色段数。在树查询的时候要注意排重的技巧，详见代码
*/
#include <stdio.h>
#include <string.h>
#include <algorithm>
#include <iostream>
using namespace std;
const int maxn=100005;
int fa[maxn],dep[maxn],siz[maxn],son[maxn],num[maxn],top[maxn];
int n,p,z,a[maxn],Hash[maxn];
int ll[maxn*4],rr[maxn*4],tree[maxn*4],flag[maxn*4];
int head[maxn],ip;

void init()
{
memset(head,-1,sizeof(head));
ip=0;
}

struct note
{
int v,next;
} edge[maxn*2];

void addedge(int u,int v)
{
edge[ip].v=v,edge[ip].next=head[u],head[u]=ip++;
}

void dfs(int u,int pre)
{
son[u]=0,siz[u]=1,dep[u]=dep[pre]+1,fa[u]=pre;
for(int i=head[u]; i!=-1; i=edge[i].next)
{
int v=edge[i].v;
if(v==pre)continue;
dfs(v,u);
siz[u]+=siz[v];
if(siz[son[u]]<siz[v])
son[u]=v;
}
// printf("%d fa dep son siz %d %d %d %d\n",u,fa[u],dep[u],son[u],siz[u]);
}
void init_que(int u ,int tp)
{
num[u]=++z,top[u]=tp,Hash[z]=u;
if(son[u])
{
init_que(son[u],tp);
}
for(int i=head[u]; i!=-1; i=edge[i].next)
{
int v=edge[i].v;
if(v==fa[u]||v==son[u])continue;
init_que(v,v);
}
//printf("%d num top %d %d\n",u,num[u],top[u]);
}
void push_down(int root)
{
if(flag[root])
{
tree[root<<1]=tree[root<<1|1]=1;
ll[root<<1]=ll[root<<1|1]=rr[root<<1]=rr[root<<1|1]=ll[root];
flag[root<<1]=flag[root<<1|1]=1;
flag[root]=0;
}
}

void push_up(int root)
{
if(ll[root<<1|1]==rr[root<<1])
tree[root]=tree[root<<1]+tree[root<<1|1]-1;
else
tree[root]=tree[root<<1]+tree[root<<1|1];
ll[root]=ll[root<<1];
rr[root]=rr[root<<1|1];
}

void build(int root,int l,int r)
{
flag[root]=0;
if(l==r)
{
tree[root]=1;
ll[root]=rr[root]=a[Hash[l]];
return;
}
int mid=(l+r)>>1;
build(root<<1,l,mid);
build(root<<1|1,mid+1,r);
push_up(root);
}
void update(int root,int l,int r,int x,int y,int z)
{
if(l>y||r<x)return;
if(x<=l&&r<=y)
{
flag[root]=1;
tree[root]=1;
ll[root]=z;
rr[root]=z;
return;
}
push_down(root);
int mid=(l+r)>>1;
update(root<<1,l,mid,x,y,z);
update(root<<1|1,mid+1,r,x,y,z);
push_up(root);
}
int query(int root,int l,int r,int x,int y)
{
if(x<=l&&r<=y)
{
return tree[root];
}
push_down(root);
int sum=0;
int mid=(l+r)>>1;
if(y<=mid)
return query(root<<1,l,mid,x,y);
else if(x>mid)
return query(root<<1|1,mid+1,r,x,y);
else
{
int sum=query(root<<1|1,mid+1,r,x,y)+query(root<<1,l,mid,x,y);
if(ll[root<<1|1]==rr[root<<1])
sum--;
return sum;
}
}

int query1(int root,int l,int r,int loc)
{
if(l==r)
{
return ll[root];
}
push_down(root);
int mid=(l+r)>>1;
if(loc<=mid)
return query1(root<<1,l,mid,loc);
else
return query1(root<<1|1,mid+1,r,loc);
}
int main()
{
//  freopen("data.txt","r",stdin);
while(~scanf("%d%d",&n,&p))
{
for(int i=1; i<=n; i++)
{
scanf("%d",&a[i]);
}
init();
for(int i=1; i<n; i++)
{
int x,y;
scanf("%d%d",&x,&y);
addedge(x,y);
addedge(y,x);
}
z=0,siz[0]=0,dep[0]=0;
dfs(1,0);
init_que(1,1);
build(1,1,z);
while(p--)
{
char s[10];
int x,y,zz;
scanf("%s",s);
if(s[0]=='Q')
{
scanf("%d%d",&x,&y);
int f1=top[x],f2=top[y],sum=0;
while(f1!=f2)
{
if(dep[f1]<dep[f2])
{
swap(f1,f2);
swap(x,y);
}
sum+=query(1,1,z,num[f1],num[x]);
if(query1(1,1,z,num[f1])==query1(1,1,z,num[fa[f1]]))
{
sum--;
}
x=fa[f1],f1=top[x];
}
if(dep[x]>dep[y])
swap(x,y);
sum+=query(1,1,z,num[x],num[y]);
printf("%d\n",sum);
}
else
{
scanf("%d%d%d",&x,&y,&zz);
int f1=top[x],f2=top[y];
while(f1!=f2)
{
if(dep[f1]<dep[f2])
{
swap(f1,f2);
swap(x,y);
}
update(1,1,z,num[f1],num[x],zz);
x=fa[f1],f1=top[x];
}
if(dep[x]>dep[y])
swap(x,y);
update(1,1,z,num[x],num[y],zz);
}
}
}
return 0;
}