COT2 - Count on a tree II

COT2 - Count on a tree II

题目描述

• 给定 \(n\) 个结点的树，每个结点有一种颜色。
• \(m\) 次询问，每次询问给出 \(u,v\)，回答 \(u,v\) 之间的路径上的结点的不同颜色数。

数据范围

• \(1\le n\le 4\times 10^4\)，\(1\le m\le 10^5\)，颜色是不超过 \(2 \times 10^9\) 的非负整数。

Solution:

树上莫队练手题。

首先我们将这颗树的欧拉序求出来，然后我们就把这个问题转成了一个序列上的问题。

对于 \(x,y\) 的关系我们分类讨论一下：

（假设 \(dep_x<dep_y\)）
假如 \(x,y\) 在一条链上，那么这段区间在欧拉序上就对应 \([ed_x,st_y]\)。直接统计就好了。

但是如果 \(x,y\) 不在一条链上，那么我们可以在欧拉序上找到 \([st_x,st_y]\) 这个区间。然后我们会发现一个很好的性质：在这个区间上，所有不在路径 (x,y) 上的点都已经被访问过两次了，而所有在路径 (x,y) 上的节点除了 \(lca\) 以外都只访问过一次。所以我们可以维护每个点 \(u\) 访问次数 \(use_u\) 的奇偶性来决定一次操作是将节点 \(u\) 插入或删除。

对于答案统计:

我们只需要维护每次插入或删除一个点过后是否有颜色个数的变化就好了。

一些细节：

由于本题欧拉序和点之间的映射关系较为复杂，所以要特别注意数组维护的是什么，下标的意义是什么，不要混淆了。

Code:

#include<bits/stdc++.h>
const int N=8e4+5;
const int M=1e5+5;
const int lg=17;
using namespace std;
int n,m,tmp,tot,S;
int st[N],dep[N],ed[N],a[N],b[N];
int ba[N],rid[N],f[N][lg+5],use[N],col[N];
int ans[M];
vector<int> E[N];
void dfs(int x,int fa)
{
    st[x]=++tot;rid[tot]=x;
    f[x][0]=fa;dep[x]=dep[fa]+1;
    for(int i=1;i<=lg;i++)f[x][i]=f[f[x][i-1]][i-1];
    for(auto y : E[x])if(y!=fa)dfs(y,x);
    ed[x]=++tot;rid[tot]=x;
}
int LCA(int x,int y)
{
    if(dep[x]<dep[y])swap(x,y);
    for(int i=lg;i>=0;i--)if(dep[f[x][i]]>=dep[y])x=f[x][i];
    for(int i=lg;i>=0;i--)if(f[x][i]!=f[y][i])x=f[x][i],y=f[y][i];
    return x==y ? x : f[x][0];
}
inline void add(int x){++ba[a[x]];tmp+=(ba[a[x]]==1);}
inline void del(int x){--ba[a[x]];tmp-=(ba[a[x]]==0);}
inline void calc(int x){(use[x]^=1) ? add(x) : del(x);}
struct task{
    int l,r,L,R,id,lca;
    bool operator <(const task &t)const{
        return L==t.L ? (L&1 ? r<t.r : t.r<r) : l<t.l;
    }
}q[M];
void work()
{
    cin>>n>>m;S=n*2/sqrt(m*2/3);
    for(int i=1;i<=n;i++)scanf("%d",&a[i]),b[i]=a[i];
    sort(b+1,b+1+n);
    for(int i=1;i<=n;i++)a[i]=lower_bound(b+1,b+1+n,a[i])-b;
    for(int i=1,x,y;i<n;i++)
    {
        scanf("%d%d",&x,&y);
        E[x].push_back(y);E[y].push_back(x);
    }
    dfs(1,0);
    for(int i=1,x,y;i<=m;i++)
    {
        scanf("%d%d",&x,&y);
        if(st[x]>st[y])swap(x,y);
        int lca=LCA(x,y),l,r,L,R;
        if(lca==x){l=st[x],r=st[y];lca=0;}
        else {l=ed[x],r=st[y];}
        L=l/S,R=r/S;
        q[i]={l,r,L,R,i,lca};
    }
    sort(q+1,q+1+m);
    int l=1,r=0;
    for(int i=1;i<=m;i++)
    {
		while (q[i].l<l)calc(rid[--l]);
		while (r<q[i].r)calc(rid[++r]);
		while (l<q[i].l)calc(rid[l++]);
		while (q[i].r<r)calc(rid[r--]);
		if(q[i].lca)calc(q[i].lca);
		ans[q[i].id]=tmp;
		if(q[i].lca)calc(q[i].lca);
    }
    for(int i=1;i<=m;i++)printf("%d\n",ans[i]);
}
int main()
{
    //freopen("Count on a tree II.in","r",stdin);freopen("Count on a tree II.out","w",stdout);
    work();
    return 0;
}