树上数数 题解

Posted 2022-04-26 little-aztl

一、题目：

题目链接

个人私题，请注意版权保护。

二、思路：

树上莫队经典例题。其莫队的思想也很经典。

首先考虑在树上怎样莫队。我们很熟悉在序列上进行莫队，那么我们考虑将树转化成序列。那么什么序列能将子树转化为一个区间呢？当然是dfs序了！

然后回忆HH的项链那道经典题。它问的是“一段区间所包含的不同颜色的数量”，我们开桶cnt[i]保存i颜色出现的数量。而这道题我们问的是“出现次数大于等于K的颜色数量”，cnt数组还是要有的，还需要什么?我们很自然地想到用某个数据结构sum高效计算出现次数大于等于i的颜色数量。

那么怎么高效维护这个数据结构呢？首先很容易想到的是log级数据结构，比如权值线段树或权值树状数组。sum[i]表示出现次数为i的颜色数量，比如对于当前扫到的颜色col，每次在更新cnt[col]的时候记原先的cnt[col]为former,更新后的cnt[col]为now，将sum[former]减一，将sum[now]加一。查询时只需查询[k,n]的区间和即可。这是可以通过本题的。时间复杂度\(O(n\sqrt n logn\)).

但是上述做法无论编程复杂度和时间复杂度都较高。我们考虑优化。更改sum的意义。sum[i]表示出现次数大于等于i的颜色数量。冷静分析。当我们准备将cnt[col]加一时，只有sum[cnt[col]+1]发生了变化，因为大于等于cnt[col]的颜色数量没有变，而大于等于cnt[col]+1的颜色数量却增加了一个。所以我们只需将sum[cnt[col]+1]++,然后将cnt[col]本身++即可。类似地，当我们准备将cnt[col]减一时，只有sum[cnt[col]]发生了变化，sum[cnt[col]-1]及其他的sum值均没有改变。时间复杂度\(O(n\sqrt n)\).

三、代码：

60分做法：

#include<iostream>
#include<cstdio>
#include<cstring>
#include<vector>
#include<algorithm>

#define LL long long
#define FILEIN(s) freopen(s".in","r",stdin)
#define FILEOUT(s) freopen(s".out","w",stdout)
#define mem(s,v) memset(s,v,sizeof(s))

using namespace std;
inline LL read(void)
    LL x=0,f=1;char ch=getchar();
    while(ch<'0'||ch>'9')if(ch=='-')f=-1;ch=getchar();
    while(ch>='0'&&ch<='9')x=x*10+ch-'0';ch=getchar();
    return f*x;


const int maxn=5005;

int n,m;

int col[maxn],b[maxn],TOT;

int head[maxn],tot;

int cnt[maxn][maxn];

struct Edge
    int y,nxt;
    Edge()
    Edge(int _y,int _nxt):y(_y),nxt(_nxt)
e[maxn<<1];

struct Query
    int v,k,ans;
q[100005];

vector<int>query[maxn];

inline void connect(int x,int y)
    e[++tot]=Edge(y,head[x]);
    head[x]=tot;


void dfs(int x,int fa)
    for(register int i=head[x];i;i=e[i].nxt)
        int y=e[i].y;
        if(y==fa)continue;
        dfs(y,x);
        for(register int j=1;j<=TOT;++j)
            cnt[x][j]+=cnt[y][j];
        
    
    cnt[x][col[x]]++;
    for(register int i=0;i<(int)query[x].size();++i)
        int cur=query[x][i];
        for(register int col=1;col<=TOT;++col)
            if(cnt[x][col]>=q[cur].k)q[cur].ans++;
        
    


int main()
//  FILEIN("count");FILEOUT("count");
    n=read();m=read();
    for(register int i=1;i<=n;++i)
        col[i]=read();
        b[++TOT]=col[i];
    
    sort(b+1,b+TOT+1);
    TOT=unique(b+1,b+TOT+1)-b-1;
    for(register int i=1;i<=n;++i)
        col[i]=lower_bound(b+1,b+TOT+1,col[i])-b;
    
    for(register int i=1;i<n;++i)
        int x=read(),y=read();
        connect(x,y);connect(y,x);
    
    for(register int i=1;i<=m;++i)
        q[i].v=read();q[i].k=read();q[i].ans=0;
        query[q[i].v].push_back(i);
    
    dfs(1,0);
    for(register int i=1;i<=m;++i)
        printf("%d\n",q[i].ans);
    
    return 0;

满分做法：

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<cmath>

#define LL long long
#define FILEIN(s) freopen(s".in","r",stdin)
#define FILEOUT(s) freopen(s".out","w",stdout)
#define mem(s,v) memset(s,v,sizeof(s))

using namespace std;
inline LL read(void)
    LL x=0,f=1;char ch=getchar();
    while(ch<'0'||ch>'9')if(ch=='-')f=-1;ch=getchar();
    while(ch>='0'&&ch<='9')x=x*10+ch-'0';ch=getchar();
    return f*x;


const int maxn=100005;

int n,m;

int head[maxn],tot,block;
int col[maxn],siz[maxn],dfn[maxn],Cnt,pos[maxn];

int sum[maxn],cnt[maxn];

struct Edge
    int y,nxt;
    Edge()
    Edge(int _y,int _nxt):y(_y),nxt(_nxt)
e[maxn<<1];

struct Query
    int l,r,k,ans,id;
q[maxn];

inline bool cmp1(Query a,Query b)
    if(a.l/block!=b.l/block)return a.l/block<b.l/block;
    return a.r<b.r;


inline bool cmp2(Query a,Query b)
    return a.id<b.id;


inline void connect(int x,int y)
    e[++tot]=Edge(y,head[x]);
    head[x]=tot;


void dfs(int x,int fa)
    dfn[++Cnt]=x;
    pos[x]=Cnt;
    siz[x]=1;
    for(register int i=head[x];i;i=e[i].nxt)
        int y=e[i].y;
        if(y==fa)continue;
        dfs(y,x);
        siz[x]+=siz[y];
    


inline void del(int x)
    --sum[cnt[col[dfn[x]]]--];


inline void add(int x)
    ++sum[++cnt[col[dfn[x]]]];


int main()
    n=read();m=read();
    block=sqrt(n);
    for(register int i=1;i<=n;++i)
        col[i]=read();
    
    for(register int i=1;i<n;++i)
        int x=read(),y=read();
        connect(x,y);connect(y,x);
    
    dfs(1,0);
    for(register int i=1;i<=m;++i)
        q[i].id=i;
        int v=read();
        q[i].k=read();
        q[i].l=pos[v];
        q[i].r=q[i].l+siz[v]-1;
        q[i].ans=0;
    
    sort(q+1,q+m+1,cmp1);
    int l=1,r=0;
    for(register int i=1;i<=m;++i)
        if(r<q[i].r)while(r!=q[i].r)add(++r);
        if(r>q[i].r)while(r!=q[i].r)del(r--);
        if(l<q[i].l)while(l!=q[i].l)del(l++);
        if(l>q[i].l)while(l!=q[i].l)add(--l);
        q[i].ans=sum[q[i].k];
    
    sort(q+1,q+m+1,cmp2);
    for(register int i=1;i<=m;++i)
        printf("%d\n",q[i].ans);
    
    return 0;