【AGC034E】Complete Compress

题目

题目链接：https://atcoder.jp/contests/agc034/tasks/agc034_e
给你一颗 \(n\) 个节点的树，并用二进制串告诉你哪些节点上有棋子（恰好一颗）。
可以进行若干次操作，每次操作可以将两颗距离至少为 \(2\) 的棋子向中间移动一步。
问能否通过若干次操作使得所有的棋子都在一个点上，如果能，输出最小操作次数，如果不能，输出 \(-1\) 。
\(2 \leq n\leq 2000\)。

思路

考虑枚举最终集合的点 \(rt\)。显然所有点到达 \(rt\) 所需要的操作次数为以 \(rt\) 为根的树时

\[\frac{\sum_{i=1}^{n} \mathrm{dep_i}}{2} \]

我们只需要 check 是否有方案使得所有点到达 \(rt\) 即可。
考虑树形 dp。设 \(f_i\) 表示所有点要到达 \(rt\) 时，\(i\) 子树内的点最多能操作多少次。
假设当前处理到 \(x\) 的子树，我们可以把一个 \(x\) 子树内的点 \(y\) 拆成 \(\mathrm{dep}_y-\mathrm{dep}_x\) 个点。那么最多的操作次数等于每次在剩余的点中选择两个不同的删去，最多能操作的次数。
设 \(\mathrm{sum}_x\) 表示所有拆出来的点的数量，\(\mathrm{maxs}_x\) 表示 \(x\) 所有儿子中 \(\mathrm{sum}_y+\mathrm{siz}_y\) 的最大值。
如果 \(2\mathrm{maxs}_x\leq \mathrm{sum_x}\)，我们不难构造出一种方案使得最终最多剩余一个点（每次删除剩余数量最多的两个点）。所以此时

\[f_x=\lfloor\frac{\mathrm{sum}_x}{2}\rfloor \]

如果 \(2\mathrm{maxs}_x> \mathrm{sum_x}\)，那么我们可以操作的次数就是 \(\mathrm{sum}_x-\mathrm{maxs}_x+f_v\)，其中 \(v\) 是 \(\mathrm{sums}_x\) 对应的子节点。
但是因为 \(v\) 子树内是需要自己匹配的，其最好情况下的匹配数量才能达到上式，而其内部匹配会导致 \(f_x\) 有一个上界 \(\lceil\frac{\mathrm{sum}_x}{2}\rceil\)。所以此时

\[f_x=\min(\mathrm{sum}_x-\mathrm{maxs}_x+f_v,\lceil\frac{\mathrm{sum}_x}{2}\rceil) \]

如果最后 \(2f_{rt}=\frac{\sum_{i=1}^{n} \mathrm{dep_i}}{2}\)，那么 \(rt\) 有解。
时间复杂度 \(O(n^2)\)。

代码

#include <bits/stdc++.h>
using namespace std;

const int N=2010,Inf=1e9;
int n,ans,deps,tot,f[N],son[N],dep[N],maxs[N],sum[N],siz[N],head[N];
bool flag[N];

struct edge
{
	int next,to;
}e[N*2];

void add(int from,int to)
{
	e[++tot]=(edge){head[from],to};
	head[from]=tot;
}

void dfs(int x,int fa)
{
	dep[x]=dep[fa]+1;
	if (flag[x]) deps+=dep[x]-1;
	siz[x]=flag[x]; maxs[x]=sum[x]=son[x]=0;
	for (int i=head[x];~i;i=e[i].next)
	{
		int v=e[i].to;
		if (v!=fa)
		{
			dfs(v,x);
			siz[x]+=siz[v]; sum[x]+=sum[v]+siz[v];
			maxs[x]=max(maxs[x],sum[v]+siz[v]);
			if (sum[v]+siz[v]>sum[son[x]]+siz[son[x]]) son[x]=v;
		}
	}
	if (maxs[x]*2<=sum[x]) f[x]=sum[x]/2;
		else f[x]=min(sum[x]-maxs[x]+f[son[x]],sum[x]-sum[x]/2);
}

int main()
{
	memset(head,-1,sizeof(head));
	scanf("%d",&n);
	for (int i=1;i<=n;i++)
		scanf("%1d",&flag[i]);
	for (int i=1,x,y;i<n;i++)
	{
		scanf("%d%d",&x,&y);
		add(x,y); add(y,x);
	}
	ans=Inf;
	for (int i=1;i<=n;i++)
	{
		tot=deps=0; dfs(i,0);
		if (f[i]*2==deps) ans=min(ans,f[i]);
	}
	if (ans==Inf) printf("-1");
		else printf("%d",ans);
	return 0;
}