BZOJ 2243 染色 | 树链剖分模板题进阶版

这道题呢_{就是个带区间修改的树链剖分}

如何区间修改？跟树链剖分的区间询问一个道理，再加上线段树的区间修改就好了。

这道题要注意的是，无论是线段树上还是原树上，把两个区间的信息合并的时候，要注意中间相邻两个颜色是否相同。

这代码好长啊啊啊啊

幸好一次过了不然我估计永远也De不出来

#include <cstdio>

#include <cstring>

#include <algorithm>

using namespace std;

typedef long long ll;

#define space putchar(' ')

#define enter putchar('\n')

template <class T>

void read(T &x){

    char c;

    bool op = 0;

    while(c = getchar(), c < '0' || c > '9')

        if(c == '-') op = 1;

    x = c - '0';

    while(c = getchar(), c >= '0' && c <= '9')

        x = x * 10 + c - '0';

    if(op) x = -x;

}

template <class T>

void write(T x){

    if(x < 0) putchar('-'), x = -x;

    if(x >= 10) write(x / 10);

    putchar('0' + x % 10);

}

char readchar(){

    char c;

    while(c = getchar(), c < 'A' || c > 'Z');

    return c;

}

const int N = 100005;

int n, m;

int ecnt, adj[N], nxt[2*N], go[2*N];

int tot, pos[N], idx[N], fa[N], son[N], sze[N], top[N], dep[N], val[N];

int le[4*N], ri[4*N], data[4*N], lazy[4*N];

void add(int u, int v){

    go[++ecnt] = v;

    nxt[ecnt] = adj[u];

    adj[u] = ecnt;

}

void pushup(int k){

    data[k] = data[k << 1] + data[k << 1 | 1];

    if(ri[k << 1] == le[k << 1 | 1]) data[k]--;

    le[k] = le[k << 1], ri[k] = ri[k << 1 | 1];

}

void pushdown(int k){

    if(lazy[k] == -1) return;

    lazy[k << 1] = lazy[k << 1 | 1] = lazy[k];

    le[k << 1] = le[k << 1 | 1] = lazy[k];

    ri[k << 1] = ri[k << 1 | 1] = lazy[k];

    data[k << 1] = data[k << 1 | 1] = 1;

    lazy[k] = -1;

}

void build(int k, int l, int r){

    lazy[k] = -1;

    if(l == r) return (void)(data[k] = 1, le[k] = ri[k] = val[idx[l]]);

    int mid = (l + r) >> 1;

    build(k << 1, l, mid);

    build(k << 1 | 1, mid + 1, r);

    pushup(k);

}

void change(int k, int l, int r, int ql, int qr, int x){

    if(ql <= l && qr >= r) return (void)(data[k] = 1, le[k] = ri[k] = lazy[k] = x);

    pushdown(k);

    int mid = (l + r) >> 1;

    if(ql <= mid) change(k << 1, l, mid, ql, qr, x);

    if(qr > mid) change(k << 1 | 1, mid + 1, r, ql, qr, x);

    pushup(k);

}

int query(int k, int l, int r, int ql, int qr){

    if(ql <= l && qr >= r) return data[k];

    pushdown(k);

    int mid = (l + r) >> 1;

    if(qr <= mid) return query(k << 1, l, mid, ql, qr);

    if(ql > mid) return query(k << 1 | 1, mid + 1, r, ql, qr);

    return query(k << 1, l, mid, ql, qr) + query(k << 1 | 1, mid + 1, r, ql, qr) - (ri[k << 1] == le[k << 1 | 1]);

}

int getcol(int k, int l, int r, int p){

    if(lazy[k] != -1) return lazy[k];

    if(l == r) return le[k];

    int mid = (l + r) >> 1;

    if(pos[p] <= mid) return getcol(k << 1, l, mid, p);

    else return getcol(k << 1 | 1, mid + 1, r, p);

}

void path_change(int u, int v, int x){

    if(top[u] == top[v]){

        if(pos[u] > pos[v]) swap(u, v);

        change(1, 1, n, pos[u], pos[v], x);

        return;

    }

    if(dep[top[u]] > dep[top[v]]) swap(u, v);

    change(1, 1, n, pos[top[v]], pos[v], x);

    path_change(u, fa[top[v]], x);

}

int path_query(int u, int v){

    if(top[u] == top[v]){

        if(pos[u] > pos[v]) swap(u, v);

        return query(1, 1, n, pos[u], pos[v]);

    }

    if(dep[top[u]] > dep[top[v]]) swap(u, v);

    int same = (getcol(1, 1, n, top[v]) == getcol(1, 1, n, fa[top[v]]));

    return path_query(fa[top[v]], u) + query(1, 1, n, pos[top[v]], pos[v]) - same;

}

void init(){

    static int que[N], qr;

    que[qr = 1] = 1;

    for(int ql = 1, u; ql <= qr; ql++){

        u = que[ql], sze[u] = 1;

        for(int e = adj[u], v; e; e = nxt[e])

            if(v = go[e], v != fa[u])

                fa[v] = u, dep[v] = dep[u] + 1, que[++qr] = v;

    }

    for(int ql = qr, u; ql; ql--){

        u = que[ql];

        sze[fa[u]] += sze[u];

        if(sze[u] >= sze[son[fa[u]]]) son[fa[u]] = u;

    }

    for(int ql = 1, u; ql <= qr; ql++)

        if(!top[u = que[ql]])

            for(int v = u; v; v = son[v])

                idx[++tot] = v, pos[v] = tot, top[v] = u;

    build(1, 1, n);

}

int main(){

    read(n), read(m);

    for(int i = 1; i <= n; i++)

        read(val[i]);

    for(int i = 1, u, v; i < n; i++)

        read(u), read(v), add(u, v), add(v, u);

    init();

    char op;

    int a, b, c;

    while(m--){

        op = readchar(), read(a), read(b);

        if(op == 'Q') write(path_query(a, b)), enter;

        else read(c), path_change(a, b, c);

    }

    return 0;

}