[kuangbin]树链剖分 D - 染色

https://vjudge.net/contest/251031#problem/D
https://blog.csdn.net/kirito_acmer/article/details/51201994
树链剖分加线段树染色维护
我的难点在于线段树的维护 lc,rc,nc,lazy

lc rc 主要是用来看线段树分开的区间进行连续询问（就是合并的时候）在分界点处颜色是否一样（我一直wa就是这个方面少考虑了一个地方）

nc是这个区间上有几个颜色

lazy标记记录的事这个区间完全被某个颜色覆盖了——加快速度

先来看看代码

一系列准备的值：

#include <iostream>
#include <cstdio>
#include <string.h>
#include <algorithm>
#include <cmath>
#define lson rt<<1,left,mid
#define rson rt<<1|1,mid+1,right
#define ls rt<<1
#define rs rt<<1|1
#define mid ((left + right) >> 1)
using namespace std;
typedef long long ll;
const int maxn = 1e5 + 1e3;
int n,m,p;
int V[maxn];
int LC,RC;//LCA向上跳的时候维护
//邻接表
struct node{
    int to,pre;
}e[maxn << 1];
int id[maxn],cnt;
//线段树
int lc[maxn << 2],rc[maxn << 2],lazy[maxn << 2],nc[maxn << 2];
//dfs1
int siz[maxn],dep[maxn],fa[maxn],son[maxn];
//dfs2
int top[maxn],num_id[maxn],id_num[maxn];
int tot;

 这里面有个LC，RC用到的时候我再说

然后初始化加边操作

void init()
{
    memset(id,-1,sizeof(id));
    memset(son,0,sizeof(son));
    cnt = tot = 0;
}
void add(int from,int to)
{
    e[cnt].to = to;
    e[cnt].pre = id[from];
    id[from] = cnt++;
}
void dfs1(int now,int f,int depth)
{
    siz[now] = 1;
    fa[now] = f;
    dep[now] = depth;

    for(int i = id[now];~i;i = e[i].pre)
    {
        int to = e[i].to;
        if(to != f)
        {
            dfs1(to,now,depth+1);
            siz[now] += siz[to];
            if(siz[to] > siz[son[now]])
                son[now] = to;
        }
    }
}
void dfs2(int now,int rt)
{
    top[now] = rt;
    num_id[now] = ++tot;
    id_num[tot] = now;

    if(!son[now]) return;

    dfs2(son[now],rt);

    for(int i = id[now];~i;i = e[i].pre)
    {
        int to = e[i].to;
        if(to != son[now] && to != fa[now])
        {
            dfs2(to,to);
        }
    }
}

 关于向上更新的操作针对于lc，rc，还有nc，很好理解，这里要注意内部边界颜色颜色相同的时候

void pup(int rt)
{
    lc[rt] = lc[ls];
    rc[rt] = rc[rs];
    nc[rt] = nc[ls] + nc[rs];
    if(rc[ls] == lc[rs])
        --nc[rt];
}

 基础建树操作

void build(int rt,int left,int right)
{
    lazy[rt] = 0;
    nc[rt] = 0;
    if(left == right)
    {

        //cout<<left<<" : "<<" "<<id_num[left] <<" "<<V[id_num[left]]<<endl;
        lc[rt] = rc[rt] = V[id_num[left]];
        nc[rt] = 1;
        return;
    }
    build(lson);
    build(rson);
    pup(rt);
}

 向下更新操作，针对lazy标记的意义没有啥难度

void pdown(int rt,int left,int right)
{
    if(lazy[rt])
    {
        int lt = lazy[rt];
        lc[ls] = rc[ls] = lc[rs] = rc[rs] = lc[rt];
        nc[ls] = nc[rs] = 1;
        lazy[rs] = lt;
        lazy[ls] = lt;
        lazy[rt] = 0;
    }
}

 重头戏之LCA+线段树更新操作

不在一条链上常规更新，在一条链上常规更新，主意好区间范围updata_lca没什么难度

对于updata函数,找到最大包围区间更新lc，rc，nc，lazy很容易，向下的更新，递归的区间操作，向上的维护，都还是基础的线段树的考察

void updata(int rt,int left,int right,int l,int r,int k)
{
    if(l <= left && right <= r)
    {
        //cout<<left<<" "<<right<<endl;
        lc[rt] = rc[rt] = k;
        nc[rt] = 1;
        lazy[rt] = 1;
        return;
    }

    pdown(rt,left,right);

    if(l <= mid)
        updata(lson,l,r,k);
    if(r > mid)
        updata(rson,l,r,k);

    pup(rt);
}
void updata_lca(int x,int y,int k)
{
    while(top[x] != top[y])
    {
        if(dep[top[x]] < dep[top[y]])swap(x,y);
        updata(1,1,tot,num_id[top[x]],num_id[x],k);
        x = fa[top[x]];
    }
    if(dep[x] < dep[y])swap(x,y);
    updata(1,1,tot,num_id[y],num_id[x],k);
}

 重头戏之LCA + 线段树查询操作

先看query_lca函数这时候我们要用到LC,RC,prex,prey这四个变量了，因为LCA借助top数组进行链链跳，把大区间分开了，所以我们在合并大区间的时候

很容易理解prex，prey分别记录的是上一条链的左头颜色和右头颜色

LC，RC呢就是Query后记录的当前区间左头又

当没有到一条链的时候我们能保证topx是深度深的，所以对于连接的时候只需要判断prex ？= Rc 即可

但是到了一条链了，你就得想这个点是从哪里并入这条链的，两种情况顶端和非顶端非顶端的情况不会交换判断prex 和 RC

顶端的时候交换了就要判断prey ？= LC的关系了，但是没什么prex和prey肯定有一个是-1没啥

这个地方我倒是想明白了，但是还解释不出来，还是动手画一画这两种情况模拟一下就好了

int query(int rt,int left,int right,int l,int r)
{
    int res = 0;
    if(left == l)LC = lc[rt];
    if(right == r)RC = rc[rt];
    if(l <= left && right <= r)
    {
        return nc[rt];
    }

    pdown(rt,left,right);

    if(l <= mid)res += query(lson,l,r);
    if(r > mid)res += query(rson,l,r);
    if(l <= mid && r > mid)
    {
        if(rc[ls] == lc[rs])res--;
    }
    return res;
}
int querty_lca(int x,int y)
{
    int res = 0;
    int prex = -1,prey = -1;
    while(top[x] != top[y])
    {
        if(dep[top[x]] < dep[top[y]]){
                swap(x,y);
                swap(prex,prey);
        }
        res += query(1,1,tot,num_id[top[x]],num_id[x]);
        if(prex == RC)res--;
        prex = LC;
        x = fa[top[x]];
    }
    if(dep[x] < dep[y])
    {
        swap(x,y);
        swap(prex,prey);
    }
    res += query(1,1,tot,num_id[y],num_id[x]);
    if(prex == RC) res--;
    if(prey == LC) res--;
    return res;
}

 好了这个题就结束了

int main()
{
    while(~scanf("%d%d",&n,&m))
    {
        init();
        for(int i=1;i<=n;++i)
        {
            scanf("%d",&V[i]);
        }
        int from,to;
        for(int i=1;i<=n-1;i++)
        {
            scanf("%d%d",&from,&to);
            add(from,to);
            add(to,from);
        }
        dfs1(1,0,1);
        dfs2(1,1);
        build(1,1,tot);
        //cout<<lc[8]<<" "<<lc[9]<<" "<<lc[5]<<" "<<lc[12]<<" "<<lc[13]<<" "<<lc[7]<<endl;
        char op[10];
        int x,y,z;
        for(int i=1;i<=m;++i)
        {
            scanf("%s",op);
            if(op[0]=='Q'){
                scanf("%d%d",&x,&y);
                printf("%d\n",querty_lca(x,y));
            }
            else
            {
                scanf("%d%d%d",&x,&y,&z);
                updata_lca(x,y,z);
            }
        }
    }
    return 0;
}