bzoj2243 SDOI2011 染色 树链剖分

树剖裸题，复习一波。。一开始还看错题目了以为要维护一段区间的颜色数量。。还以为要上主席树。。

#include<cstdio>#include<cstring>#include<algorithm>#include<iostream>#include<queue>#define fo(i,a,b) for(int i=a;i<=b;i++)#define fd(i,a,b) for(int i=a;i>=b;i--)using namespace std;const int N=5e5+5;int n,m,cnt;int val[N],head[N],next[N],go[N],fa[N],pos[N];int a[N],tot;int siz[N],son[N],top[N],dep[N];struct node{    int l,r,sum,lc,rc,lazy;}t[N*10];inline void add(int x,int y){    go[++tot]=y;    next[tot]=head[x];    head[x]=tot;}inline void dfs1(int x,int f){    siz[x]=1;    for(int i=head[x];i;i=next[i])    {        int v=go[i];        if (v!=f)        {            dep[v]=dep[x]+1;            fa[v]=x;            dfs1(v,x);            siz[x]+=siz[v];            if (!son[x]||siz[son[x]]<siz[v])            son[x]=v;        }    }}int tim;inline void dfs2(int x,int chain){    int k=0;    pos[x]=++tim;    top[x]=chain;    for(int i=head[x];i;i=next[i])    {        int v=go[i];        if (v!=fa[x])        {            if (siz[v]>siz[k])            k=v;        }    }    if (!k)return ;    dfs2(k,chain);    for(int i=head[x];i;i=next[i])    {        int v=go[i];        if (v!=fa[x]&&v!=k)        dfs2(v,v);    }}inline void build(int x,int l,int r){    t[x].l=l,t[x].r=r;    if (l==r)return ;    int mid=(l+r)>>1;    build(x*2,l,mid);    build(x*2+1,mid+1,r);}inline void pushlazy(int x){    if (t[x].lazy==0||t[x].l==t[x].r)return;    int v=t[x].lazy;    t[x].lazy=0;    t[x<<1].sum=t[x<<1|1].sum=1;    t[x<<1].lc=t[x<<1].rc=t[x<<1|1].lc=t[x<<1|1].rc=t[x<<1].lazy=t[x<<1|1].lazy=v;}inline void queryc(int x,int l,int r,int v){    if (t[x].l==l&&t[x].r==r)    {        t[x].sum=1;        t[x].lc=t[x].rc=t[x].lazy=v;        return;    }    pushlazy(x);    int mid=(t[x].l+t[x].r)>>1;    if (r<=mid)queryc(x<<1,l,r,v);    else if (l>mid)queryc(x<<1|1,l,r,v);    else    {        queryc(x<<1,l,mid,v);        queryc(x<<1|1,mid+1,r,v);    }    t[x].lc=t[x<<1].lc;    t[x].rc=t[x<<1|1].rc;    t[x].sum=t[x<<1].sum+t[x<<1|1].sum;    if (t[x<<1].rc==t[x<<1|1].lc)t[x].sum--;}inline void solvec(int x,int y,int z){    while (top[x]!=top[y])    {        if (dep[top[x]]<dep[top[y]])swap(x,y);        queryc(1,pos[top[x]],pos[x],z);        x=fa[top[x]];    }    if (dep[x]>dep[y])swap(x,y);    queryc(1,pos[x],pos[y],z);}inline int queryq(int x,int l,int r){    pushlazy(x);    if (t[x].l==l&&t[x].r==r)return t[x].sum;    int mid=(t[x].l+t[x].r)>>1;    if (r<=mid)return queryq(x<<1,l,r);    else if (l>mid)return queryq(x<<1|1,l,r);    else    {        int sum=queryq(x<<1,l,mid)+queryq(x<<1|1,mid+1,r);        if (t[x<<1].rc==t[x<<1|1].lc)sum--;        return sum;    }}inline int find(int x,int pos){    pushlazy(x);    if (t[x].sum==1)return t[x].lc;    int mid=(t[x].l+t[x].r)>>1;    if (pos<=mid)return find(x<<1,pos);    else find(x<<1|1,pos);}inline int solveq(int x,int y){    int sum=0;    while (top[x]!=top[y])    {        if (dep[top[x]]<dep[top[y]])swap(x,y);        sum+=queryq(1,pos[top[x]],pos[x]);        if (find(1,pos[top[x]])==find(1,pos[fa[top[x]]]))sum--;//注意合并的时候相邻颜色相同sum-1         x=fa[top[x]];    }        if (dep[x]>dep[y])swap(x,y);    sum+=queryq(1,pos[x],pos[y]);    return sum;}int main(){    scanf("%d%d",&n,&m);    fo(i,1,n)scanf("%d",&a[i]);    fo(i,1,n-1)    {        int x,y;        scanf("%d%d",&x,&y);        add(x,y);        add(y,x);    }    dep[1]=1;    dfs1(1,0);    dfs2(1,1);    build(1,1,n);    char ch[2];    fo(i,1,n)    queryc(1,pos[i],pos[i],a[i]);    fo(i,1,m)    {        scanf("%s",ch);        if (ch[0]=='C')        {            int x,y,z;            scanf("%d%d%d",&x,&y,&z);            solvec(x,y,z);        }        else        {            int x,y;            scanf("%d%d",&x,&y);            printf("%d\n",solveq(x,y));        }    }}