bzoj 2243 [SDOI2011]染色 （线段树 + 树链剖分）

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

题解

首先处理树，使用树链剖分；线段树则维护区间颜色不同数量及左右端点的颜色。更新时若左边的右端点和右边的左端点颜色相同，则数量减1。

代码

/**************************************************************
    Problem: 2243
    User: xingsuixing
    Language: C++
    Result: Accepted
    Time:5276 ms
    Memory:31076 kb
****************************************************************/

#include<cstdio>
#include<cstring>
#include<algorithm>
#include<iostream>
using namespace std;

const int N = 200010;

struct Edge {
    int v, next;
} edge[N];
int n, m, tot = 0, prelc, prerc, root, tail = 0;
int indexx = 0, fa[N], dep[N], top[N], siz[N], son[N], in[N], out[N], a[N], head[N], color[N];

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

struct Node {
    int l, r, ls, rs;
    int sum, flag, lc, rc;
} t[500010];

void dfs1(int u, int f, int d) {
    siz[u] = 1, fa[u] = f, dep[u] = d;
    for(int i = head[u]; i; i = edge[i].next) {
        int v = edge[i].v;
        if(v == f) continue ;
        dfs1(v, u, d + 1);
        siz[u] += siz[v];
        if(son[u] == -1 || siz[v] > siz[son[u]])
            son[u] = v;
    }
}

void dfs2(int u, int tp) {
    indexx ++;
    in[u] = out[u] = indexx, top[u] = tp;
    color[indexx] = a[u];
    if(son[u] == -1) return ;
    dfs2(son[u], tp);
    for(int i = head[u]; i; i = edge[i].next) {
        int v = edge[i].v;
        if(v == fa[u] || v == son[u]) continue ;
        dfs2(v, v);
    }
    out[u] = indexx;
}

void update(int num) {
    int ls = num << 1, rs = num << 1 | 1;
    t[num].lc = t[ls].lc;
    t[num].rc = t[rs].rc;
    t[num].sum = t[ls].sum + t[rs].sum;
    if(t[ls].rc == t[rs].lc)
        t[num].sum -= 1;
}

void build(int num, int l, int r) {
    t[num].l = l, t[num].r = r;
    if(l == r) {
        t[num].sum = 1, t[num].flag = 0;
        t[num].lc = t[num].rc = color[l];
        return ;
    }
    int mid = (l + r) >> 1;
    build(num << 1, l, mid);
    build(num << 1 | 1, mid + 1, r);
    update(num);
}

void pushdown(int num) {
    int ls = num << 1, rs = num << 1 | 1;
    if(t[num].flag) {
        t[ls].sum = t[rs].sum = 1;
        t[ls].lc = t[ls].rc = t[num].flag;
        t[rs].lc = t[rs].rc = t[num].flag;
        t[rs].flag = t[ls].flag = t[num].flag;
        t[num].flag = 0;
    }
}

void modify(int num, int L, int R, int c) {
    int l = t[num].l, r = t[num].r;
    if(L <= l && r <= R) {
        t[num].flag = c;
        t[num].sum = 1;
        t[num].lc = c, t[num].rc = c; 
        return ;
    }
    pushdown(num);
    int mid = (l + r) >> 1;
    if(L <= mid)
        modify(num << 1, L, R, c);
    if(R > mid)
        modify(num << 1 | 1, L, R, c);
    update(num);
}

void delta(int u, int v, int c) {
    int f1 = top[u], f2 = top[v];
    while(f1 != f2){
        if(dep[f1] < dep[f2]) swap(f1, f2), swap(u, v);  
        modify(root, in[f1], in[u], c);
        u = fa[f1], f1 = top[u];
    }
    if(dep[u] > dep[v]) swap(u, v);
    modify(root, in[u], in[v], c);
}

int query(int num, int pos, int val, int L, int R) {
    int l = t[num].l, r = t[num].r;
    if(l == L) prelc = t[num].lc;
    if(r == R) prerc = t[num].rc;
    if(L <= l && r <= R)
        return t[num].sum;
    pushdown(num);
    int mid = (l + r) >> 1, rt = 0;
    if(L > mid)
        rt = query(num << 1 | 1, pos, val, L, R);
    else if(R <= mid)
        rt = query(num << 1, pos, val, L, R);
    else {
        rt = query(num << 1, pos, mid, L, R) + query(num << 1 | 1, mid + 1, val, L, R);
        if(t[num << 1].rc == t[num << 1 | 1].lc)
            rt -= 1;
    } 
    return rt;
}

int check(int u, int v) {
    int co1 = -1, co2 = -1;
    int f1 = top[u], f2 = top[v], sum = 0;
    while(f1 != f2) {
        if(dep[f1] < dep[f2]) swap(f1, f2), swap(u, v), swap(co1, co2);
        sum += query(1, in[f1], in[u], in[f1], in[u]);
        if(prerc == co1) sum --;
        co1 = prelc;
        u = fa[f1];
        f1 = top[u];
    }
    if(dep[u] < dep[v]) swap(u, v), swap(co1, co2);
    sum += query(1, in[v], in[u], in[v], in[u]);
    if(prerc == co1) sum --;
    if(prelc == co2) sum --;
    return sum;
}

int main(){
    memset(son, -1, sizeof(son));
    scanf("%d %d", &n, &m);
    for(int i = 1; i <= n; i ++)
        scanf("%d", &a[i]);
    for(int i = 1; i < n; i ++) {
        int u, v;
        scanf("%d %d", &u, &v);
        add(u, v), add(v, u);
    }
    dfs1(1, -1, 0); dfs2(1, 1);
    int root = 1;
    build(1, 1, n);
    char s[2];
    while(m --) {
        scanf("%s", s);
        if(s[0] == 'Q') {
            int u, v;
            scanf("%d %d", &u, &v);
            printf("%d\n", check(u, v));
        }
        else {
            int u, v, c;
            scanf("%d %d %d", &u, &v, &c);
            delta(u, v, c);
        }
    }
    return 0;
}