BZOJ 2243: [SDOI2011]染色 【树链剖分】

2243: [SDOI2011]染色

Time Limit: 20 Sec   Memory Limit: 512 MB
Submit: 4577   Solved: 1709
[ Submit][ Status][ Discuss]

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]之间。

思路：在链上跑线段树区间合并。不过有坑点，查询时沿着不同的链向上找，注意连续两条链的端点，不能单纯统计连续段数目，要去掉重复计算的。幸好test不错，要不会WA死的。

AC代码：

#include <iostream>
#include <cstdio>
#include <cstring>
#include <cmath>
#include <cstdlib>
#include <algorithm>
#include <queue>
#include <stack>
#include <map>
#include <set>
#include <vector>
#include <string>
#define INF 0x3f3f3f3f
#define eps 1e-8
#define MAXN (100000+10)
#define MAXM (300000+10)
#define Ri(a) scanf("%d", &a)
#define Rl(a) scanf("%lld", &a)
#define Rf(a) scanf("%lf", &a)
#define Rs(a) scanf("%s", a)
#define Pi(a) printf("%d\n", (a))
#define Pf(a) printf("%.2lf\n", (a))
#define Pl(a) printf("%lld\n", (a))
#define Ps(a) printf("%s\n", (a))
#define W(a) while((a)--)
#define CLR(a, b) memset(a, (b), sizeof(a))
#define MOD 1000000007
#define LL long long
#define lson o<<1, l, mid
#define rson o<<1|1, mid+1, r
#define ll o<<1
#define rr o<<1|1
#define PI acos(-1.0)
#pragma comment(linker, "/STACK:102400000,102400000")
#define fi first
#define se second
using namespace std;
struct Tree{
    int l, r, sum, lc, rc, lazy;
};
Tree tree[MAXN<<2];
void PushUp(int o){
    tree[o].sum = tree[ll].sum + tree[rr].sum;
    tree[o].lc = tree[ll].lc;
    tree[o].rc = tree[rr].rc;
    if(tree[ll].rc == tree[rr].lc)
        tree[o].sum--;
}
void PushDown(int o)
{
    if(tree[o].lazy != -1)
    {
        tree[ll].lazy = tree[rr].lazy = tree[o].lazy;
        tree[ll].lc = tree[ll].rc = tree[rr].lc = tree[rr].rc = tree[o].lazy;
        tree[ll].sum = tree[rr].sum = 1;
        tree[o].lazy = -1;
    }
}
void Build(int o, int l, int r)
{
    tree[o].l = l; tree[o].r = r;
    tree[o].lc = tree[o].rc = -1, tree[o].sum = 0; tree[o].lazy = -1;
    if(l == r)
        return ;
    int mid = (l + r) >> 1;
    Build(lson); Build(rson);
}
void Update(int o, int L, int R, int v)
{
    if(tree[o].l == L && tree[o].r == R)
    {
        tree[o].lc = tree[o].rc = v;
        tree[o].sum = 1;
        tree[o].lazy = v;
        return ;
    }
    PushDown(o);
    int mid = (tree[o].l + tree[o].r) >> 1;
    if(R <= mid) Update(ll, L, R, v);
    else if(L > mid) Update(rr, L, R, v);
    else {Update(ll, L, mid, v); Update(rr, mid+1, R, v);}
    PushUp(o);
}
int Query(int o, int L, int R)
{
    if(tree[o].l == L && tree[o].r == R)
        return tree[o].sum;
    PushDown(o);
    int mid = (tree[o].l + tree[o].r) >> 1;
    if(R <= mid) return Query(ll, L, R);
    else if(L > mid) return Query(rr, L, R);
    else return Query(ll, L, mid) + Query(rr, mid+1, R) - (tree[ll].rc == tree[rr].lc);
}
int Color(int o, int pos)
{
    if(tree[o].l == tree[o].r)
        return tree[o].lc;
    PushDown(o);
    int mid = (tree[o].l + tree[o].r) >> 1;
    if(pos <= mid) return Color(ll, pos);
    else return Color(rr, pos);
}
struct Edge{
    int from, to, next;
};
Edge edge[MAXN<<1];
int head[MAXN], edgenum;
void init(){
    edgenum = 0; CLR(head, -1);
}
void addEdge(int u, int v)
{
    Edge E = {u, v, head[u]};
    edge[edgenum] = E;
    head[u] = edgenum++;
}
int son[MAXN], num[MAXN];
int top[MAXN], pos[MAXN], id;
int dep[MAXN], pre[MAXN];
void DFS1(int u, int fa, int d)
{
    dep[u] = d; pre[u] = fa; num[u] = 1; son[u] = -1;
    for(int i = head[u]; i != -1; i = edge[i].next)
    {
        int v = edge[i].to;
        if(v == fa) continue;
        DFS1(v, u, d+1);
        num[u] += num[v];
        if(son[u] == -1 || num[son[u]] < num[v])
            son[u] = v;
    }
}
void DFS2(int u, int T)
{
    top[u] = T; pos[u] = ++id;
    if(son[u] == -1) return ;
    DFS2(son[u], T);
    for(int i = head[u]; i != -1; i = edge[i].next)
    {
        int v = edge[i].to;
        if(v == pre[u] || v == son[u]) continue;
        DFS2(v, v);
    }
}
int GetSum(int u, int v)
{
    int f1 = top[u], f2 = top[v];
    int ans = 0;
    while(f1 != f2)
    {
        if(dep[f1] < dep[f2])
        {
            swap(u, v);
            swap(f1, f2);
        }
        ans += Query(1, pos[f1], pos[u]) - (Color(1, pos[f1]) == Color(1, pos[pre[f1]]));
        u = pre[f1], f1 = top[u];
    }
    if(dep[u] > dep[v]) swap(u, v);
    return ans += Query(1, pos[u], pos[v]);
}
void Change(int u, int v, int z)
{
    int f1 = top[u], f2 = top[v];
    while(f1 != f2)
    {
        if(dep[f1] < dep[f2])
        {
            swap(u, v);
            swap(f1, f2);
        }
        Update(1, pos[f1], pos[u], z);
        u = pre[f1], f1 = top[u];
    }
    if(dep[u] > dep[v]) swap(u, v);
    Update(1, pos[u], pos[v], z);
}
int c[MAXN];
int main()
{
    int n, m;
    while(scanf("%d%d", &n, &m) != EOF)
    {
        for(int i = 1; i <= n; i++) Ri(c[i]);
        init();
        for(int i = 1; i <= n-1; i++)
        {
            int s, e;
            Ri(s), Ri(e);
            addEdge(s, e);
            addEdge(e, s);
        }
        DFS1(1, -1, 1); id = 0; DFS2(1, 1); Build(1, 1, id);
        for(int i = 1; i <= n; i++) Update(1, pos[i], pos[i], c[i]);
        W(m)
        {
            char op[5];
            Rs(op); int x, y, z;
            if(op[0] == 'C')
            {
                Ri(x); Ri(y); Ri(z);
                Change(x, y, z);
            }
            else
            {
                Ri(x); Ri(y);
                Pi(GetSum(x, y));
            }
        }
    }
    return 0;
}