HYSBZ 2243 染色

题目链接：http://www.lydsy.com/JudgeOnline/problem.php?id=2243

题意:

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

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

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

思路：

权值在点上的树链剖分+线段树成段更新+线段树成段询问。

维护颜色端的数量方法：

如果左儿子区间的最右端颜色与右儿子区间的最左端区间相同。

sum[rt] = sum[rt<<1]+sum[rt<<1|1] - 1;

否则sum[rt] = sum[rt<<1]+sum[rt<<1|1];

树链剖分部分也进行相应维护。

 

#include <iostream>
#include <cstring>
#include <string>
#include <cstdio>
#include <vector>
using namespace std;
int n, q;
#define maxn 100010
#define lson l, m, rt<<1
#define rson m+1, r, rt<<1|1
int siz[maxn], top[maxn], fa[maxn], son[maxn], dep[maxn];
int w[maxn], fw[maxn];
int A[maxn];
vector <int> mp[maxn];
int pos;
int Left, Right;
int sum[maxn<<2], cl[maxn<<2], cr[maxn<<2], col[maxn<<2];
void PushUp(int rt)
{
    if(cr[rt<<1] == cl[rt<<1|1]) sum[rt] = sum[rt<<1] + sum[rt<<1|1] - 1;
    else sum[rt] = sum[rt<<1] + sum[rt<<1|1];
    cl[rt] = cl[rt<<1];
    cr[rt] = cr[rt<<1|1];
}
void PushDown(int rt, int m)
{
    if(col[rt] >= 0)
    {
        col[rt<<1] = col[rt<<1|1] = col[rt];
        sum[rt<<1] = sum[rt<<1|1] = 1;
        cl[rt<<1] = cl[rt<<1|1] = cr[rt<<1] = cr[rt<<1|1] = cl[rt];
        col[rt] = -1;
    }
}
void build(int l, int r, int rt)
{
    col[rt] = -1;
    if(l == r)
    {
        cl[rt] = cr[rt] = A[fw[l]];
        sum[rt] = 1;
        return;
    }
    int m = (l+r)>>1;
    build(lson);
    build(rson);
    PushUp(rt);
}
void update(int L, int R, int c, int l, int r, int rt)
{
    if(L <= l && R >= r)
    {
        col[rt] = c; sum[rt] = 1; cl[rt] = c; cr[rt] = c;
        return;
    }
    PushDown(rt, r-l+1);
    int m = (l+r)>>1;
    if(L <= m) update(L, R, c, lson);
    if(R > m) update(L, R, c, rson);
    PushUp(rt);
}
int query(int L, int R, int l, int r, int rt)
{
    if(l == L) Left = cl[rt];
    if(r == R) Right = cr[rt];
    if(L <= l && R >= r)
    {
        return sum[rt];
    }
    PushDown(rt, r-l+1);
    int m = (l+r)>>1;
    int ret = 0;
    //合并线段树 这里也要修改!
    if(R <= m) return query(L, R, lson);
    else if(L > m) return query(L, R, rson);
    else
    {
        ret = query(L,R,lson) + query(L, R, rson);
        if(cl[rt<<1|1] == cr[rt<<1]) ret --;
        return ret;
    }
    //if(L <= m) ret += query(L, R, lson);
    //if(R > m) ret += query(L, R, rson);
}
int dfs1(int u, int pre, int deep)
{
    siz[u] = 1; dep[u] = deep; fa[u] = pre;
    int mmax = 0;
    for(int i = 0; i < mp[u].size(); i++)
    {
        if(mp[u][i] != pre)
        {
            int temp = dfs1(mp[u][i], u, deep+1);
            siz[u] += temp;
            if(son[u] == -1 || temp >= mmax) 
            {
                son[u] = mp[u][i];
                mmax = temp;
            }
        }
    }
    return siz[u];
}
void dfs2(int u, int val)
{
    top[u] = val;
    if(son[u] != -1)
    {
        w[u] = ++pos;
        fw[w[u]] = u;
        dfs2(son[u], val);
    }
    else if(son[u] == -1)
    {
        w[u] = ++pos;
        fw[w[u]] = u;
        return;
    }
    for(int i = 0; i < mp[u].size(); i++)
    {        
        if(mp[u][i] != son[u] && mp[u][i] != fa[u]) dfs2(mp[u][i], mp[u][i]);
    }
}
int find(int u, int v)
{
    int mark1 = -1, mark2 = -1;
    int f1 = top[u], f2 = top[v];
    int temp = 0; int precolor = -1;
    while(f1 != f2)
    {
        if(dep[f1] < dep[f2])   //规定f1更深
        {
            swap(f1, f2);
            swap(u, v);
            swap(mark1, mark2);  //交替
        }
        temp += query(w[f1], w[u], 1, pos, 1);
        if(Right == mark1) temp --;
        mark1 = Left;
        u = fa[f1]; f1 = top[u];
    }
    if(dep[u] < dep[v])  //规定u更深
    {
        swap(u, v);
        swap(mark1, mark2);
    }
    temp += query(w[v], w[u], 1, pos, 1);
    if(Right == mark1) temp--;
    if(Left == mark2) temp--;
    return temp;
}
void UPDATE(int u, int v, int color)
{
    int f1 = top[u], f2 = top[v];
    while(f1 != f2)
    {
        if(dep[f1] < dep[f2])
        {
            swap(f1, f2);
            swap(u, v);
        }
        update(w[f1], w[u], color, 1, pos, 1);
        u = fa[f1]; f1 = top[u];
    }
    if(dep[u] > dep[v]) swap(u, v);
    update(w[u], w[v], color, 1, pos, 1);
}
int main() 
{
    //freopen("in.txt", "r", stdin);
    while(~scanf("%d%d", &n, &q))
    {
        for(int i = 1; i <= n; i++) mp[i].clear(); 
        pos = 0;
        memset(son, -1, sizeof(son));
        for(int i = 1; i <= n; i++) scanf("%d", &A[i]);
        for(int i = 1; i <= n-1; i++)
        {
            int a, b; scanf("%d%d", &a, &b);
            mp[a].push_back(b);
            mp[b].push_back(a);
        }
        dfs1(1, -1, 1);
        dfs2(1, 1);
        build(1, pos, 1);
        
        char op[5];
        while(q--)
        {
            scanf("%s", op);
            int a, b, c;
            if(op[0] == 'Q')
            {
                scanf("%d%d", &a, &b);
                printf("%d\n", find(a, b));
            }
            else if(op[0] == 'C')
            {
                scanf("%d%d%d", &a, &b, &c);
                UPDATE(a, b, c);
            }
        }
    }
    return 0;
}