CF1527D MEX Tree

题目描述

给出一棵 n 个点的树，点从 0 到 n – 1 编号。定义一条路径的权值是路径上所有点编号的 \operatorname{mex}。对于每个0\le i\le n 。求出\operatorname{mex} 为 i 的路径有几条。注意，这里统计的路径需要包括至少一条边。

一个集合的 \operatorname{mex} 定义为最小的不在集合中的非负整数。

Solution

我们设 f(i) 表示 \operatorname{mex}\geq i 的情况数量。

那么 \operatorname{mex} = i 的情况数量就是 f(i+1)-f(i)

考虑计算 f(i)

f(0) 就是以 0 为根时， 0 的每个儿子的子树内部的点相互为路径的两端点。

即
\sum\limits_{x \in s}\frac{size(x)*(size(x)-1)}{2}
f(1) 即为以 0 为根时，每对兄弟子树内部的点相互为路径的两端点。

考虑计算 i \geq 2 的情况，使用双指针法，p 与 q 分别指向一条链的两端点。

设当前维护到 f(x) ，那么这条链一定是包含 [0,x-1] 的所有点 ，否则无法构成 \operatorname{mex}\geq i

然后考虑从 x 拓展到 x+1 ，详见代码

#include<bits/stdc++.h>
#define ll long long
const int N = 2e5+10;
int t,n,fa[N],cov[N];
ll sz[N],ans[N],fl,p,q;
std::vector<int> g[N];
void dfs(int x,int f){
    fa[x] = f ;
    sz[x] = 1 ; 
    for(auto p:g[x]){
        if(p==f) continue;
        dfs(p,x);
        sz[x]+=sz[p];
    }
}
void add(int x){
    if(cov[x]) return ;
    if(!cov[fa[x]]) add(fa[x]); 
    cov[x] = 1;
    if(fa[x]!=q&&fa[x]!=p){
        fl = 0 ;
        return ;
    }
    else{
         if(fa[x]==p) p=x;
         else q=x;
         if(fa[x]==0) sz[0]-=sz[x];
    }
}
int main(){
    scanf("%d",&t);
    while(t--){
        scanf("%d",&n);
        for(int i=0;i<=n;++i) g[i].clear(),ans[i]=0,cov[i]=0;
        for(int i=1;i<=n-1;++i){
            int u,v;
            scanf("%d%d",&u,&v);
            g[u].push_back(v);
            g[v].push_back(u);
        }
        dfs(0,-1);
        ll sum = 1 ;
        for(auto p:g[0]){
            ans[0]+=sz[p]*(sz[p]-1)/2;
            ans[1]+=sz[p]*sum;
            sum+=sz[p];
        }
        fl = 1 , p = q = 0; cov[0] = 1;
        for(int i=1;i<n;++i){
            add(i);
            if(!fl) break;
            ans[i] -= 1ll*sz[p]*sz[q];
            ans[i+1] = 1ll*sz[p]*sz[q];
        } 
        for(int i=0;i<=n;++i)
            printf("%lld ",ans[i]);
        puts("");
    }
    return 0;
}