P9603 [IOI 2023] 山毛榉树 题解

题目链接：P9603 [IOI 2023] 山毛榉树

刻画条件。对于一个颜色 $ c $，设排列中颜色为 $ c $ 的点分别为 $ x_1, x_2, \dots, x_k $，根据题目条件有 $ fa_{x_1}=v_0, fa_{x_2}=v_1, \dots, fa_{x_k}=v_{k-1} $，即前 $ k $ 个点都恰好有且只有一条颜色为 $ c $ 的儿子边，并且这些儿子在排列中出现的顺序和父亲出现的顺序一致。显然可以发现在序列中靠前的节点包含的颜色一定包含靠后的节点，且因为同样颜色的儿子顺序前面也大于后面，可以得到前面的节点的子树大小一定大于后面的节点的子树大小。

定义 $ S(u) $ 为 $ u $ 的出边的颜色集合，$ sz_u $ 代表 $ u $ 的子树大小，定义 $ x \preceq y $ 当且仅当 $ S(x) \subseteq S(y) $ 并且有 $ sz_{s} \le sz_{t} $，其中 $ s, t $ 分别为 $ x, y $ 的儿子，并且 $ x, y $ 到 $ s, t $ 所经过的那一条边颜色相同。这个大小比较是有传递性的，并且如果两个子树大小相等且一个小于另一个那么交换后也可以得出另一个小于这一个，因此是符合排序的要求的。此时原题目的判定可以转化为考虑一个点 $ u $ 的子树，符合条件等价于将 $ u $ 的儿子按照子树大小从小到大排序 $ son_1, son_2, \dots, son_k $，有 $ son_1 \preceq son_2 \preceq \dots \preceq son_k $，下面将证明这个判定是充要的。

必要性：根据上面论证这个小于关系符合排序性质，并且这个判定符合上面刻画的条件。

充分性：对于一个颜色 $ c $，假设两个连向父亲的颜色为 $ c $ 的边的父亲分别为 $ fa_i, fa_j $，其中 $ i<j $，因为 $ fa_j \preceq fa_i $，所以有 $ sz_{s} \le sz_{t} $，其中 $ s, t $ 分别是那两个儿子，所以这些儿子也按照子树大小顺序降序出现。如果有两个儿子的子树大小相同，那么可以靠父亲节点的顺序，因此符合条件。

综上我们完成了对条件的转换，启发式合并即可做到 $ O(nlog^2n) $。

代码：

#include<bits/stdc++.h>
#define pii pair<int,int>
using namespace std;
vector<pii > e[200005];
set<pii > s[200005];
int sz[200005];
vector<int> ans;
bool check(int x,int y)
{
	if(e[x].size()>e[y].size())return 0;
	for(auto u:e[x])
	{
		auto it=lower_bound(e[y].begin(),e[y].end(),make_pair(u.first,-1));
		if(it==e[y].end()||(*it).first!=u.first||sz[u.second]>sz[(*it).second])return 0;
	}
	return 1;
}
bool merge(int u,int v)
{
	if(s[u].size()<s[v].size())swap(s[u],s[v]);
	for(auto x:s[v])
	{
		auto it=s[u].lower_bound(x);
		if(it!=s[u].end()&&!check(x.second,(*it).second))return 0;
		if(it!=s[u].begin())
		{
			it--;
			if(!check((*it).second,x.second))return 0;
		}
		s[u].insert(x);
	}
	s[v].clear();
	return 1;
}
void dfs(int u)
{
	sz[u]=ans[u]=1;
	for(auto v:e[u])
	{
		dfs(v.second);
		sz[u]+=sz[v.second];
		if(!ans[v.second])ans[u]=0;
	}
	for(int i=1;i<e[u].size();i++)
	{
		if(e[u][i].first==e[u][i-1].first)ans[u]=0;
	}
	if(!ans[u])return ;
	s[u].insert({sz[u],u});
	for(auto v:e[u])
	{
		if(!merge(u,v.second))
		{
			ans[u]=0;
			return ;
		}
	}
}
vector<int> beechtree(int n,int m,vector<int> fa,vector<int> col)
{
	for(int i=0;i<n;i++)
	{
		e[i].clear();
		s[i].clear();
	}
	ans.assign(n,0);
	for(int i=1;i<n;i++)
	{
		e[fa[i]].push_back({col[i],i});
	}
	for(int i=0;i<n;i++)
	{
		sort(e[i].begin(),e[i].end());
	}
	dfs(0);
	return ans;
}