HYSBZ - 2243 树链剖分 + 线段树 处理树上颜色段数

Posted 2021-03-13 zzidun-pavo

用线段树处理颜色段数

记录区间内的颜色段数,区间右端点的颜色,区间右端点的颜色.

int tr[maxn<<2], lc[maxn<<2], rc[maxn<<2];

懒标记,记录区间是否被覆盖

int lazy[maxn<<2];

合并的方法是这样,对于某一区间

? 如果 左区间的右端点颜色 == 右区间的左端点

? 那么 这左右区间合并,左区间的最右边一段和右区间最左边一段颜色是连续的,那么区间的颜色段数为 左区间颜色段数+右区间颜色段数 - 1.

? 否则 区间的颜色段数为 左区间颜色段数+右区间颜色段数.

void pushup(int rt){
    lc[rt] = lc[rt*2];
    rc[rt] = rc[rt*2+1];
    tr[rt] = tr[rt*2] + tr[rt*2+1];
    if(rc[rt*2] == lc[rt*2+1]) tr[rt]--;
}

当区间的lazy不为0时, 他的左右区间都被覆盖,而且颜色段数都是1.

void pushdown(int rt){
    if(lazy[rt]){
        tr[rt*2] = tr[rt*2+1] = 1;
        lazy[rt*2] = lazy[rt*2+1] = lazy[rt];
        lc[rt*2] = lc[rt*2+1] = lc[rt];
        rc[rt*2] = rc[rt*2+1] = rc[rt];
        lazy[rt] = 0;
    }
}

建树和区间修改的操作如下

建树时,结点的初始值时字典序为l的点的颜色,rnk是字典序到点编号的映射.

void build(int rt, int l, int r){
    lazy[rt] = 0;
    if(l == r){
        tr[rt] = 1;
        lc[rt] = rc[rt] = col[rnk[l]];
        return;
    }
    int m = (l+r)/2;
    build(rt*2, l, m);
    build(rt*2+1, m+1, r);
    pushup(rt);
}

void change(int rt, int l, int r, int ql, int qr, int val){
    if(ql <= l && r <= qr){
        lazy[rt] = 1;
        tr[rt] = 1;
        lc[rt] = rc[rt] = val;
        return;
    }
    pushdown(rt);
    int m = (l+r)/2;
    if(ql <= m) change(rt*2, l, m, ql, qr, val);
    if(m < qr) change(rt*2+1, m+1, r, ql, qr, val);
    pushup(rt);
}

查询的时候要注意,如果查询的区间被分成两段,那么也要判断中间是否连续:

int query(int rt, int l, int r, int ql, int qr){
    if(ql <= l && r <= qr) return tr[rt];
    pushdown(rt);
    int m = (l+r)/2;
    if(qr <= m) return query(rt*2, l, m, ql, qr);
    else if(ql > m) return query(rt*2+1, m+1, r, ql, qr);
    else{
        int ret = query(rt*2, l, m, ql, qr) + query(rt*2+1, m+1, r, ql, qr);
        if(rc[rt*2] == lc[rt*2+1]) ret--;
        return ret;
    }
}

树链剖分

先是定义存图

vector<int> G[maxn];

然后是每个点的深度,大小,重儿子,父亲.

int dep[maxn], siz[maxn], son[maxn], fa[maxn];

然后是每个点所在的重链的顶点的编号,dfs顺序和点编号的映射,tot记录顺序的编号.

int top[maxn], tid[maxn], rnk[maxn], tot;

第一次dfs,确定点之间的父子关系,大小,深度,重儿子:

void dfs1(int u, int f){
    fa[u] = f, son[u] = 0;
    dep[u] = dep[f]+1, siz[u] = 1;
    for(int i = 0; i < G[u].size(); ++i){
        int v = G[u][i];
        if(v == fa[u]) continue;
        dfs1(v, u);
        siz[u] += siz[v];
        if(siz[v] > siz[son[u]])
            son[u] = v;
    }
}

第二次dfs,给每个结点安排重链编号,dfs序号:

void dfs2(int u, int tp){
    tot++;
    tid[u] = tot, rnk[tot] = u;
    top[u] = tp;
    if(son[u])
        dfs2(son[u], tp);
    for(int i = 0; i < G[u].size(); ++i){
        int v = G[u][i];
        if(v == son[u] || v == fa[u]) continue;
        dfs2(v, v);
    }
}

修改操作

和普通树链剖分一致.

void changeTree(int u, int v, int val){
    while(top[u] != top[v]){
        if(dep[top[u]] < dep[top[v]]) swap(u, v);
        change(1, 1, tot, tid[top[u]], tid[u], val);
        u = fa[top[u]];
    }
    if(dep[u] > dep[v]) swap(u, v);
    change(1, 1, tot, tid[u], tid[v], val);
}

查询操作

和普通的树链剖分相比,多了一步.

要判断u所在重链的顶点top[u]和顶点的父节点fa[top[u]]颜色是否相等,是的话ret要减一(代表这一条链上的颜色和上一条的连续).

int queryTree(int u, int v){
    int ret = 0;
    while(top[u] != top[v]){
        if(dep[top[u]] < dep[top[v]]) swap(u, v);
        ret += query(1, 1, tot, tid[top[u]], tid[u]);
        if(query2(1, 1, tot, tid[top[u]]) == query2(1, 1, tot, tid[fa[top[u]]])) --ret;
        u = fa[top[u]];
    }
    if(dep[u] > dep[v]) swap(u, v);
    ret += query(1, 1, tot, tid[u], tid[v]);
    return ret;
}

完整代码

#include<bits/stdc++.h>

using namespace std;

const int maxn = 100100;

int n, m;
vector<int> G[maxn];
int col[maxn];

int dep[maxn], siz[maxn], son[maxn], fa[maxn];
int top[maxn], tid[maxn], rnk[maxn], tot;

void dfs1(int u, int f){
    fa[u] = f, son[u] = 0;
    dep[u] = dep[f]+1, siz[u] = 1;
    for(int i = 0; i < G[u].size(); ++i){
        int v = G[u][i];
        if(v == fa[u]) continue;
        dfs1(v, u);
        siz[u] += siz[v];
        if(siz[v] > siz[son[u]])
            son[u] = v;
    }
}

void dfs2(int u, int tp){
    tot++;
    tid[u] = tot, rnk[tot] = u;
    top[u] = tp;
    if(son[u])
        dfs2(son[u], tp);
    for(int i = 0; i < G[u].size(); ++i){
        int v = G[u][i];
        if(v == son[u] || v == fa[u]) continue;
        dfs2(v, v);
    }
}

int tr[maxn<<2], lc[maxn<<2], rc[maxn<<2];
int lazy[maxn<<2];

void pushup(int rt){
    lc[rt] = lc[rt*2];
    rc[rt] = rc[rt*2+1];
    tr[rt] = tr[rt*2] + tr[rt*2+1];
    if(rc[rt*2] == lc[rt*2+1]) tr[rt]--;
}

void pushdown(int rt){
    if(lazy[rt]){
        tr[rt*2] = tr[rt*2+1] = 1;
        lazy[rt*2] = lazy[rt*2+1] = lazy[rt];
        lc[rt*2] = lc[rt*2+1] = lc[rt];
        rc[rt*2] = rc[rt*2+1] = rc[rt];
        lazy[rt] = 0;
    }
}

void build(int rt, int l, int r){
    lazy[rt] = 0;
    if(l == r){
        tr[rt] = 1;
        lc[rt] = rc[rt] = col[rnk[l]];
        return;
    }
    int m = (l+r)/2;
    build(rt*2, l, m);
    build(rt*2+1, m+1, r);
    pushup(rt);
}

void change(int rt, int l, int r, int ql, int qr, int val){
    if(ql <= l && r <= qr){
        lazy[rt] = 1;
        tr[rt] = 1;
        lc[rt] = rc[rt] = val;
        return;
    }
    pushdown(rt);
    int m = (l+r)/2;
    if(ql <= m) change(rt*2, l, m, ql, qr, val);
    if(m < qr) change(rt*2+1, m+1, r, ql, qr, val);
    pushup(rt);
}

int query(int rt, int l, int r, int ql, int qr){
    if(ql <= l && r <= qr) return tr[rt];
    pushdown(rt);
    int m = (l+r)/2;
    if(qr <= m) return query(rt*2, l, m, ql, qr);
    else if(ql > m) return query(rt*2+1, m+1, r, ql, qr);
    else{
        int ret = query(rt*2, l, m, ql, qr) + query(rt*2+1, m+1, r, ql, qr);
        if(rc[rt*2] == lc[rt*2+1]) ret--;
        return ret;
    }
}

int query2( int rt, int l, int r, int x){
    if(l == r) return lc[rt];
    pushdown(rt);
    int m = (l+r)/2;
    if(x <= m) return query2(rt*2, l, m, x);
    else return query2(rt*2+1, m+1, r, x);
}

void changeTree(int u, int v, int val){
    while(top[u] != top[v]){
        if(dep[top[u]] < dep[top[v]]) swap(u, v);
        change(1, 1, tot, tid[top[u]], tid[u], val);
        u = fa[top[u]];
    }
    if(dep[u] > dep[v]) swap(u, v);
    change(1, 1, tot, tid[u], tid[v], val);
}

int queryTree(int u, int v){
    int ret = 0;
    while(top[u] != top[v]){
        if(dep[top[u]] < dep[top[v]]) swap(u, v);
        ret += query(1, 1, tot, tid[top[u]], tid[u]);
        if(query2(1, 1, tot, tid[top[u]]) == query2(1, 1, tot, tid[fa[top[u]]])) --ret;
        u = fa[top[u]];
    }
    if(dep[u] > dep[v]) swap(u, v);
    ret += query(1, 1, tot, tid[u], tid[v]);
    return ret;
}

int main()
{
    scanf("%d%d", &n,&m);
        for(int i = 1; i <= n; ++i) G[i].clear();
        for(int i = 1; i <= n; ++i) scanf("%d", &col[i]);
        for(int i = 1; i < n; ++i){
            int u, v;
            scanf("%d%d", &u, &v);
            G[u].push_back(v);
            G[v].push_back(u);
        }
        dfs1(1, 0);
        tot = 0;
        dfs2(1, 1);

        build(1, 1, tot);

        while(m--){
            char opt;
            scanf(" %c", &opt);
            if(opt == 'Q'){
                int u, v;
                scanf("%d%d", &u, &v);
                printf("%d
", queryTree(u, v));
            }else if(opt == 'C'){
                int u, v, val;
                scanf("%d%d%d", &u, &v, &val);
                changeTree(u, v, val);
            }
        }
}