【AGC034E】Complete Compress

题目

题目链接：https://atcoder.jp/contests/agc034/tasks/agc034_e
给你一颗 (n) 个节点的树，并用二进制串告诉你哪些节点上有棋子（恰好一颗）。
可以进行若干次操作，每次操作可以将两颗距离至少为 (2) 的棋子向中间移动一步。
问能否通过若干次操作使得所有的棋子都在一个点上，如果能，输出最小操作次数，如果不能，输出 (-1) 。
(2 leq nleq 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) 的最大值。
如果 (2mathrm{maxs}_xleq mathrm{sum_x})，我们不难构造出一种方案使得最终最多剩余一个点（每次删除剩余数量最多的两个点）。所以此时

[f_x=lfloorfrac{mathrm{sum}_x}{2} floor ]

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

[f_x=min(mathrm{sum}_x-mathrm{maxs}_x+f_v,lceilfrac{mathrm{sum}_x}{2} ceil) ]

如果最后 (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;
}