bzoj4999 This Problem Is Too Simple！树链剖分+线段树

Description

---

给您一颗树，每个节点有个初始值。
现在支持以下两种操作：
1. C i x(0<=x<2^31) 表示将i节点的值改为x。
2. Q i j x(0<=x<2^31) 表示询问i节点到j节点的路径上有多少个值为x的节点。

Solution

---

最开始想到3个log的做法，显然是过不了的
注意到多种颜色很烦，那么我们可以考虑每种颜色单独建一棵线段树，然后树链剖分求和即可

如果n和q很小的话可以考虑树上莫队

Code

---

#include <stdio.h>
#include <string.h>
#include <map>
#define rep(i,st,ed) for (int i=st;i<=ed;++i)

const int N=100005;
const int E=200005;

struct edge {int y,next;} e[E];
struct treeNode {int l,r,sum;} t[N*61];

std:: map <int,int> map;

int bl[N],pos[N],dep[N],fa[N],size[N];
int root[N*3],c[N],tot,color;
int ls[N],edCnt;

int read() {
    int x=0,v=1; char ch=getchar();
    for (;ch<'0'||ch>'9';v=(ch=='-')?(-1):(v),ch=getchar());
    for (;ch<='9'&&ch>='0';x=x*10+ch-'0',ch=getchar());
    return x*v;
}

void add_edge(int x,int y) {
    e[++edCnt]=(edge) {y,ls[x]}; ls[x]=edCnt;
    e[++edCnt]=(edge) {x,ls[y]}; ls[y]=edCnt;
}

void dfs1(int now) {
    size[now]=1;
    for (int i=ls[now];i;i=e[i].next) {
        if (e[i].y==fa[now]) continue;
        dep[e[i].y]=dep[now]+1;
        fa[e[i].y]=now;
        dfs1(e[i].y);
        size[now]+=size[e[i].y];
    }
}

void dfs2(int now,int up) {
    pos[now]=++pos[0]; bl[now]=up;
    int mx=0;
    for (int i=ls[now];i;i=e[i].next) {
        if (e[i].y==fa[now]) continue;
        if (size[e[i].y]>size[mx]) mx=e[i].y;
    }
    if (!mx) return ;
    dfs2(mx,up);
    for (int i=ls[now];i;i=e[i].next) {
        if (e[i].y==fa[now]||e[i].y==mx) continue;
        dfs2(e[i].y,e[i].y);
    }
}

void modify(int &now,int tl,int tr,int x,int v) {
    if (!now) now=++tot;
    t[now].sum+=v;
    if (tl==tr) return ;
    int mid=(tl+tr)>>1;
    if (x<=mid) modify(t[now].l,tl,mid,x,v);
    else if (x>mid) modify(t[now].r,mid+1,tr,x,v);
}

int query(int now,int tl,int tr,int l,int r) {
    if (!now) return 0;
    if (tl==l&&tr==r) return t[now].sum;
    int mid=(tl+tr)>>1;
    if (r<=mid) return query(t[now].l,tl,mid,l,r);
    if (l>mid) return query(t[now].r,mid+1,tr,l,r);
    return query(t[now].l,tl,mid,l,mid)+query(t[now].r,mid+1,tr,mid+1,r);
}

void solve(int x,int y,int col) {
    int ans=0;
    while (bl[x]!=bl[y]) {
        if (dep[bl[x]]<dep[bl[y]]) std:: swap(x,y);
        ans+=query(root[map[col]],1,pos[0],pos[bl[x]],pos[x]);
        x=fa[bl[x]];
    }
    if (pos[x]>pos[y]) std:: swap(x,y);
    ans+=query(root[map[col]],1,pos[0],pos[x],pos[y]);
    printf("%d\n", ans);
}

int main(void) {
    int n=read(),T=read();
    rep(i,1,n) {
        c[i]=read();
        if (!map[c[i]]) {
            map[c[i]]=++color;
        }
    }
    rep(i,2,n) add_edge(read(),read());
    dep[1]=1; dfs1(1); dfs2(1,1);
    rep(i,1,n) modify(root[map[c[i]]],1,n,pos[i],1);
    for (;T--;) {
        char opt[2]; scanf("%s",opt);
        if (opt[0]=='Q') {
            int x=read(),y=read(),col=read();
            solve(x,y,col);
        } else if (opt[0]=='C') {
            int x=read(),col=read();
            modify(root[map[c[x]]],1,n,pos[x],-1);
            c[x]=col;
            if (!map[c[x]]) {
                map[c[x]]=++color;
            }
            modify(root[map[c[x]]],1,n,pos[x],1);
        }
    }
    return 0;
}