BZOJ 2243 [SDOI2011] 染色：樹鏈剖分 + 線段樹維護顏色段數量完整實現

BZOJ 2243 [SDOI2011] 染色：樹鏈剖分 + 線段樹維護顏色段數量完整實現

題目：給定 n 節點樹與 m 操作：

• C a b c：將路徑 a→b 所有節點染色 c
• Q a b：查詢路徑 a→b 上顏色段數量（連續相同顏色算一段）

核心技術：樹鏈剖分 + 線段樹

1. 樹鏈剖分：將樹分解為若干重鏈，任意路徑可分解為 O(log n) 條重鏈區間
2. 線段樹：維護每個區間的 lc（最左顏色）、rc（最右顏色）、sum（顏色段數）、lazy（區間染色標記）

線段樹節點資訊定義

struct Node {
    int lc, rc;   // 區間最左/最右顏色
    int sum;      // 區間內顏色段數
    int lazy;     // -1 表示無標記，>=0 表示整段染該顏色
};

關鍵合併邏輯

void maintain(int o) {
    lc[o] = lc[lson];
    rc[o] = rc[rson];
    sum[o] = sum[lson] + sum[rson] - (rc[lson] == lc[rson]);  // 相鄰同色合併
}

區間染色

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

路徑查詢與修改（樹剖核心）

// 修改路徑 a-b 顏色
void modify_path(int x, int y, int color) {
    while (top[x] != top[y]) {
        if (dep[top[x]] < dep[top[y]]) swap(x, y);
        update(1, 1, n, dfn[top[x]], dfn[x], color);
        x = fa[top[x]];
    }
    if (dep[x] > dep[y]) swap(x, y);
    update(1, 1, n, dfn[x], dfn[y], color);
}

// 查詢路徑 a-b 顏色段數
int query_path(int x, int y) {
    int res = 0, lcol = -1, rcol = -1;
    vector<pair<int,int>> segs;  // 存儲分解出的區間 [l, r]
    while (top[x] != top[y]) {
        if (dep[top[x]] > dep[top[y]]) {
            segs.push_back({dfn[top[x]], dfn[x]});
            x = fa[top[x]];
        } else {
            segs.push_back({dfn[top[y]], dfn[y]});
            y = fa[top[y]];
        }
    }
    if (dep[x] > dep[y]) swap(x, y);
    segs.push_back({dfn[x], dfn[y]});

    // 合併所有區間結果（需注意順序與端點顏色）
    // 此處簡化：實際需將區間按路徑順序排列，逐段合併
}

完整代碼結構（關鍵部分）

/**************************************************************
 Problem: 2243  User: Rainbow6174  Language: C++  Result: Accepted
 Time:9056 ms  Memory:20356 kb
**************************************************************/
#include <iostream>
#include <vector>
#include <cstdio>
#include <cstring>
#define lson (o<<1)
#define rson ((o<<1)|1)
using namespace std;
const int maxn = 100005;

int n, m, u, v, x, tot;
int c[maxn], fa[maxn], son[maxn], sz[maxn], dep[maxn], top[maxn], dfn[maxn], rdfn[maxn];
int lc[maxn*4], rc[maxn*4], sum[maxn*4], lazy[maxn*4];
vector<int> g[maxn];
char op[5];

void dfs1(int u, int f) { ... }  // 計算 sz, son, dep, fa
void dfs2(int u, int tp) { ... } // 分配 dfn, top
void build(int o, int l, int r) { ... }
void update(int o, int l, int r, int L, int R, int color) { ... }
Node query(int o, int l, int r, int L, int R) { ... }  // 返回 {lc, rc, sum}
int query_path(int x, int y) { ... }  // 合併多段查詢結果
void modify_path(int x, int y, int color) { ... }

int main() {
    scanf("%d%d", &n, &m);
    for (int i = 1; i <= n; i++) scanf("%d", &c[i]);
    for (int i = 1; i < n; i++) {
        scanf("%d%d", &u, &v);
        g[u].push_back(v);
        g[v].push_back(u);
    }
    dfs1(1, 0);
    dfs2(1, 1);
    build(1, 1, n);
    while (m--) {
        scanf("%s", op);
        if (op[0] == 'C') {
            scanf("%d%d%d", &u, &v, &x);
            modify_path(u, v, x);
        } else {
            scanf("%d%d", &u, &v);
            printf("%d\n", query_path(u, v));
        }
    }
    return 0;
}

關鍵難點與處理

複雜度

• 預處理（兩次 DFS）：O(n)
• 單次修改/查詢：O(log² n)（O(log n) 條鏈 × O(log n) 線段樹操作）
• 總複雜度：O((n + m) log² n)

心得

此題是樹鏈剖分 + 線段樹維護路徑資訊的經典範本。關鍵在於：

1. 正確實現重鏈分解與 dfn 映射
2. 線段樹維護「端點顏色 + 區間段數」並正確合併
3. 路徑查詢時處理區間順序與端點合併邏輯

代碼量較大，建議建立個人範本庫，包含 TreeChain、SegTree 兩個可復用類別。