【BZOJ 2243】【SDOI 2011】染色【树链剖分】

学会树链剖分后的一道练手题。>_<

Description

给定一棵有n个节点的无根树和m个操作，操作有2类：
1、将节点a到节点b路径上所有点都染成颜色c；
2、询问节点a到节点b路径上的颜色段数量（连续相同颜色被认为是同一段），如“112221”由3段组成：“11”、“222”和“1”。
请你写一个程序依次完成这m个操作。

题解

若是你还不会树链剖分，请看这里

　　本题可以直接上 树链剖分+线段树。
　　在维护线段树时，要保存sum（这段区间的段数），setv（是否有区间更新），lc，rc（左右两边的节点颜色）。更新节点时，只需注意左儿子的rc和右儿子的lc是否相同即可。
　　在剖分树链时，也要注意每条链的左右颜色，若相同，则ans-1。

代码

#include<cstdio>
#include<cstring>
#include<algorithm>

#define N 100010
using namespace std;

struct node{int to, next;}e[N*2];
int head[N],m;

int tim,n;
int dep[N],fa[N],son[N],size[N],top[N],tid[N],Rank[N],a[N];
int lco,rco;

void add_edge(int from,int to)
{
    e[++m].next = head[from];
    head[from] = m;
    e[m].to = to;
}

void dfs1(int v,int pa,int d)
{
    dep[v] = d; fa[v] = pa; size[v] = 1;
    for(int i = head[v];i;i=e[i].next)
        if(e[i].to != pa)
        {
            dfs1(e[i].to,v,d+1);
            size[v] += size[e[i].to];
            if(son[v] == -1 || size[e[i].to] > size[son[v]])
                son[v] = e[i].to;
        }
}

void dfs2(int v,int tp)
{
    top[v] = tp; tid[v] = ++tim;
    Rank[tid[v]] = v;
    if(son[v] == -1) return;
    dfs2(son[v],tp);
    for(int i = head[v];i;i=e[i].next)
        if(e[i].to != fa[v] && e[i].to != son[v])
            dfs2(e[i].to,e[i].to);
}

#define lson o << 1
#define rson o << 1 | 1
int sum[N<<2],setv[N<<2],lc[N<<2],rc[N<<2];

void pushup(int o) {
    sum[o] = sum[lson] + sum[rson];
    lc[o] = lc[lson];
    rc[o] = rc[rson];
    if(rc[lson] == lc[rson]) sum[o]--;
}

void pushdown(int o)
{
    if(setv[o] != -1) {
        setv[lson] = setv[rson] = setv[o];
        lc[lson] = rc[lson] = setv[o];
        lc[rson] = rc[rson] = setv[o];
        sum[lson] = sum[rson] = 1;
        setv[o] = -1;
    }
}

void build(int o,int l,int r)
{
    setv[o] = -1;
    if(l == r) {
        sum[o] = 1;
        lc[o] = rc[o] = a[Rank[l]];
        return;
    }
    int mid = (l+r)>>1;
    build(lson,l,mid); build(rson,mid+1,r);
    pushup(o);
}

void update(int o,int l,int r,int ll,int rr,int v)
{
    if(ll <= l && rr >= r) {
        setv[o] = v;
        sum[o] = 1;
        lc[o] = rc[o] = v;
        return;
    }
    pushdown(o);
    int mid = (l+r)>>1;
    if(ll <= mid) update(lson,l,mid,ll,rr,v);
    if(rr > mid) update(rson,mid+1,r,ll,rr,v);
    pushup(o);
}

void change(int x,int y,int v)
{
    while(top[x] != top[y]) {
        if(dep[top[x]] < dep[top[y]]) swap(x,y);
        update(1,1,n,tid[top[x]],tid[x],v);
        x = fa[top[x]];
    }
    if(dep[x] > dep[y]) swap(x,y);
    update(1,1,n,tid[x],tid[y],v);
}

int query(int o,int l,int r,int ll,int rr)
{
    if(l == ll) lco = lc[o];
    if(r == rr) rco = rc[o];
    if(ll <= l && rr >= r) return sum[o];
    pushdown(o);
    int mid = (l+r)>>1,ans = 0;
    if(ll <= mid) ans += query(lson,l,mid,ll,rr);
    if(rr > mid) ans += query(rson,mid+1,r,ll,rr);
    if(ll <= mid && rr > mid && rc[lson] == lc[rson]) ans--;
    pushup(o);
    return ans;
}

int Query(int x,int y)
{
    int ans = 0,ans1 = -1,ans2 = -1;
    while(top[x] != top[y]) {
        if(dep[top[x]] < dep[top[y]]) {swap(x,y);swap(ans1,ans2);}
        ans += query(1,1,n,tid[top[x]],tid[x]);
        //printf("%d %d %d\n",lco,rco,ans1);
        //printf("1--%d\n",ans);
        if(ans1 == rco) ans--;
        //printf("2--%d\n",ans);
        ans1 = lco;
        x = fa[top[x]];
    }
    if(dep[x] > dep[y]) {swap(x,y);swap(ans1,ans2);}
    ans += query(1,1,n,tid[x],tid[y]);
    //printf("%d %d\n",lco,rco);
    if(ans1 == lco) ans--;
    if(ans2 == rco) ans--;
    return ans;
}   

int main()
{
    char opt[3];
    int q,u,v,ll,rr;
    scanf("%d%d",&n,&q);
    for(int i = 1;i <= n;i++) scanf("%d",&a[i]);

    memset(son,-1,sizeof(son));
    memset(head,0,sizeof(head));
    tim = m = 0;

    for(int i = 1;i < n;i++) {
        scanf("%d%d",&u,&v);
        add_edge(u,v); add_edge(v,u);
    }
    dfs1(1,0,0);
    dfs2(1,1);
    build(1,1,n);

    while(q--) {
        scanf("%s%d%d",opt,&ll,&rr);
        if(opt[0] == 'Q') printf("%d\n",Query(ll,rr));
        else { scanf("%d",&v); change(ll,rr,v); }
    }
    return 0;
}