[Luogu] 染色

https://www.luogu.org/problemnew/show/P2486

qizha

为什么会wa

#include <cstdio>
#include <cmath>
#include <iostream>
#include <cstring>

using std:: cin;
using std:: cout;
using std:: endl;
using std:: swap;
using std:: string;

#define gc getchar()

const int N = 1e5 + 10;

//...数据 
int n, m, data[N];
//...图 
int now = 1, head[N];
struct Node {int u, v, nxt;} G[N << 1];
//...树剖 
int tim;
int fa[N], deep[N], size[N], topp[N], son[N], tree[N], bef[N];
//...线段树 
int L[N << 2], R[N << 2], F[N << 2], l_col[N << 2], r_col[N << 2], tot[N << 2];
//...查询操作
int Lc, Rc, ans; 

inline int read() {
    int x = 0, f = 1; char c = gc;
    while(c < '0' || c > '9') {if(c == '-') f = -1; c = gc;}
    while(c >= '0' && c <= '9') x = x * 10 + c - '0', c = gc;
    return x * f;
}

inline void Add(int u, int v) {G[now].v = v; G[now].nxt = head[u]; head[u] = now ++;}

void Dfs_find_son(int u, int f_, int dep) {
    fa[u] = f_;
    deep[u] = dep;
    size[u] = 1;
    for(int i = head[u]; ~ i; i = G[i].nxt) {
        int v = G[i].v;
        if(v != f_) {
            Dfs_find_son(v, u, dep + 1);
            size[u] += size[v];
            if(size[v] > size[son[u]]) son[u] = v;
        }
    }
}

void Dfs_to_un(int u, int tp) {
    topp[u] = tp;
    tree[u] = ++ tim;
    bef[tim] = u;
    if(! son[u]) return ;
    Dfs_to_un(son[u], tp);
    for(int i = head[u]; ~ i; i = G[i].nxt) {
        int v = G[i].v;
        if(v != fa[u] && v != son[u]) Dfs_to_un(v, v);
    }
}

#define lson jd << 1
#define rson jd << 1 | 1

void Up(int jd) {
    l_col[jd] = l_col[lson]; r_col[jd] = r_col[rson];
    tot[jd] = tot[lson] + tot[rson];
    if(r_col[lson] == l_col[rson]) tot[jd] --;
}

void Build_tree(int l, int r, int jd) {
    L[jd] = l; R[jd] = r; F[jd] = -1;
    if(l == r) {
        l_col[jd] = r_col[jd] = data[bef[l]];
        tot[jd] = 1;
        return ;
    }
    int mid = (l + r) >> 1;
    Build_tree(l, mid, lson);
    Build_tree(mid + 1, r, rson);
    Up(jd);
}

void Down_f(int jd) {
    F[lson] = F[rson] = F[jd];
    tot[lson] = tot[rson] = 1;
    l_col[lson] = r_col[lson] = l_col[rson] = r_col[rson] = F[jd];
    F[jd] = -1;
    return ;
}

void Sec_A(int l, int r, int jd, int x, int y) {
    if(x == L[jd]) Lc = l_col[jd];
    if(y == R[jd]) Rc = r_col[jd];
    if(x <= l && r <= y) {ans += tot[jd]; return ;}
    if(F[jd] != -1) Down_f(jd);
    int mid = (l + r) >> 1;
    if(x <= mid) Sec_A(l, mid, lson, x, y);
    if(y > mid)  Sec_A(mid + 1, r, rson, x, y);
    if(x <= mid && y > mid && r_col[lson] == l_col[rson]) ans --;
}

inline int Sec_A_imp(int x, int y) {
    int tp1 = topp[x], tp2 = topp[y], ret = 0, prex = -1, prey = -1;
    while(tp1 != tp2) {
        if(deep[tp1] < deep[tp2]) swap(tp1, tp2), swap(prex, prey), swap(x, y);
        ans = 0;
        Sec_A(1, n, 1, tree[tp1], tree[x]);
        ret += ans;
        if(Rc == prex) ret --;
        prex = Lc;
        x = fa[tp1];
        tp1 = topp[x];
    }
    if(deep[x] < deep[y]) swap(x, y), swap(prex, prey);
    ans = 0;
    Sec_A(1, n, 1, tree[y], tree[x]);
    ret += ans;
    if(Lc == prey) ret --;
    if(Rc == prex) ret --;
    return ret;
}

void Sec_G(int l, int r, int jd, int x, int y, int how) {
    if(x <= l && r <= y) {
        F[jd] = how; tot[jd] = 1; l_col[jd] = r_col[jd] = how;
        return ;
    }
    int mid = (l + r) >> 1;
    if(x <= mid) Sec_G(l, mid, lson, x, y, how);
    if(y > mid)  Sec_G(mid + 1, r, rson, x, y, how);
    Up(jd);
}

void Sec_G_imp(int x, int y, int how) {
    int tp1 = topp[x], tp2 = topp[y];
    while(tp1 != tp2) {
        if(deep[tp1] < deep[tp2]) swap(tp1, tp2), swap(x, y);
        Sec_G(1, n, 1, tree[tp1], tree[x], how);
        x = fa[tp1];
        tp1 = topp[x];
    }
    if(deep[x] < deep[y]) swap(x, y);
    Sec_G(1, n, 1, tree[y], tree[x], how);
    return ;
}

#define RR freopen("gg.in", "r", stdin)

int main() {
    //RR;
    n = read(); m = read();
    for(int i = 1; i <= n; i ++) data[i] = read();
    for(int i = 1; i <= n; i ++) head[i] = -1;
    for(int i = 1; i <  n; i ++) {
        int u = read(), v = read();
        Add(u, v); Add(v, u);
    }
    Dfs_find_son(1, 0, 1);
    Dfs_to_un(1, 1);
    Build_tree(1, n, 1);
    while(m --) {
        string s;
        cin >> s;
        if(s[0] == 'Q') {
            int x = read(), y = read();
            cout << Sec_A_imp(x, y) << endl;
        } else {
            int x = read(), y = read(), how = read();
            Sec_G_imp(x, y, how);
        }    
    }
    return 0;
}
/*
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
*/