洛谷 SP10707 COT2 - Count on a tree II(树上莫队)

链接：

https://www.luogu.org/problemnew/show/SP10707

题意：

给定一棵$n$个节点的树，每个节点都有一种颜色$a[x]$，给定一些查询$(x,y)$，问$x$到$y$的路径上共有多少种颜色；

分析：

树上莫队，即把树上路径问题转化为区间问题；
(之后想补再回来补。。。)下面代码是wa的，没发现哪错了。。

代码：

#include<algorithm>
#include<iostream>
#include<cstring>
#include<string>
#include<cstdio>
#include<vector>
#include<queue>
#include<stack>
#include<cmath>
#include<set>
#include<map>

using namespace std;

#define  inf 0x7f7f7f7f
#define  maxn 100050
#define  N 100100
#define  P 2

typedef long long ll;
typedef struct {
    int u, v, next, w;
} Edge;
Edge e[maxn];
int cnt, head[maxn];

inline void add(int u, int v) {
    e[cnt].u = u;
    e[cnt].v = v;
    //e[cnt].w=w;
    // e[cnt].f=f;
    e[cnt].next = head[u];
    head[u] = cnt++;
    e[cnt].u=v;
    e[cnt].v=u;
//    e[cnt].w=0;
//    e[cnt].f=-f;
    e[cnt].next=head[v];
    head[v]=cnt++;
}

inline void write(int x) {
    if (x < 0)
        putchar('-'), x = -x;
    if (x > 9)
        write(x / 10);
    putchar(x % 10 + '0');
}

inline int read() {
    int x = 0, f = 1;
    char c = getchar();
    while (c < '0' || c > '9') {
        if (c == '-')
            f = -1;
        c = getchar();
    }
    while (c >= '0' && c <= '9') {
        x = x * 10 + c - '0';
        c = getchar();
    }
    return x * f;
}
struct node{
    int l,r,id,lca;
}p[maxn];
int c[maxn],res,dep[maxn],num,sz,cot[maxn],n,m,pos[maxn],a[maxn],sa[maxn],f[maxn][20],st[maxn],ed[maxn],vis[maxn],ans[maxn];
bool com(node a,node b){
    if(pos[a.l]!=pos[b.l])return pos[a.l]<pos[b.l];
    else return a.r<b.r;
}
void dfs(int x,int fa){
    st[x]=++num;sa[num]=x;
    dep[x]=dep[fa]+1;
    f[x][0]=fa;
    for(int i=1;i<=16;i++){
        f[x][i]=f[f[x][i-1]][i-1];
    }
    for(int i=head[x];i!=-1;i=e[i].next){
        int v=e[i].v;
        if(v==fa)continue;
        dfs(v,x);
    }
    ed[x]=++num;sa[num]=x;
}
int lca(int a,int b){
    if(dep[b]>dep[a])swap(a,b);
    for(int i=16;i>=0;i--)
        if(dep[f[a][i]]>=dep[b])a=f[a][i] ;
    if(a==b) return a;
    for(int i=16;i>=0;i--){
        if(f[a][i]!=f[b][i])
            a=f[a][i],b=f[b][i];
    }
    return f[a][0];
}
void _del(int x){
    if(cot[a[x]]==1)res--;
    cot[a[x]]--;
}
void _add(int x){
    if(cot[a[x]]==0)res++;
    cot[a[x]]++;
}
void update(int x){
    vis[x]?_del(x):_add(x);
    vis[x]^=1;
}
int main() {
    cin>>n>>m;
    sz=sqrt(n);
    for(int i=1;i<=n;i++) {
        scanf("%d", &a[i]);
        c[i]=a[i];
    }
    sort(c+1,c+1+n);
    unique(c+1,c+1+n);
    for(int i=1;i<=n;i++)
        a[i]=lower_bound(c+1,c+1+n,a[i])-c;
    for(int i=1;i<=2*n;i++)
        pos[i]=i/sz;
    memset(head,-1,sizeof(head));
    int a,b;
    for(int i=1;i<=n-1;i++){
        scanf("%d%d",&a,&b);
        add(a,b);
    }
    dfs(1,0);
    for(int i=1;i<=m;i++){
        scanf("%d%d",&a,&b);
        p[i].id=i;
        if(st[a]>st[b])swap(a,b);
        int _lca=lca(a,b);
        if(_lca==st[a])
            p[i].lca=0, p[i].l=st[a],p[i].r=st[b];
        else p[i].lca=_lca,p[i].l=ed[a],p[i].r=st[b];
    }
    sort(p+1,p+1+m,com);
    int l=1,r=0;
    for(int i=1;i<=m;i++){
        while(l<p[i].l){
            update(sa[l]);
            l++;
        }
        while(l>p[i].l){
            l--;
            update(sa[l]);
        }
        while(r<p[i].r){
            r++;
            update(sa[r]);
        }
        while(r>p[i].r){
            update(sa[r]);
            r--;
        }
        if(p[i].lca)update(p[i].lca);
        ans[p[i].id]=res;
        if(p[i].lca)update(p[i].lca);
    }
    for(int i=1;i<=m;i++)
        printf("%d\n",ans[i]);
    return 0;
}

我们坚持一件事情，并不是因为这样做了会有效果，而是坚信，这样做是对的。
——哈维尔