CF1830D Mex Tree

Description

给定一棵 $n$ 个结点的树，点有点权，点权为 0 或 1。你需要给一种指定点权的方案，使得每一条路径的点权 $\operatorname{mex}$ 之和最大。

$n\le 2\times 10^5$，多组数据。

3s , 250MB

Solution

显然答案的上界是所有点对的贡献都是 $2$，上界为 $n(n+1)$ 。

有一种显然的构造方法，黑白交替染色，答案的下界是 $n(n+1)-n-\sum[p_i=1]$。

考虑计算贡献差 $2$ 的最小值。

而贡献不为 $2$ 的点对都是在同色的连通块中。颜色为 $c$，大小为 $k$ 的同色极大连通块的贡献为 $(c+1)k(k+1)/2$。

那么就可以 $dp$ 了。令 $f_{x,j,0/1}$ 表示在 $x$ 的子树内，颜色为 $0/1$，当前的极大连通块大小为 $j$。这是好转移的。但是状态大小是 $O(n^2)$ 的，有点愚蠢。

考虑优化掉无效状态。因为已经构造出一种黑白交替的方案，最多只会减 $2n$。而一个连通块的贡献是 $O(n^2)$ 级别的，所以 $dp$ 的第二维只需要开 $\sqrt n$。这样的时间复杂度参考树上背包，是 $O(n\sqrt n) $。

但是如果数组直接开这么大也会 GG。考虑用动态数组。转移完以后用 shrink_to_fit() 清空内存。

#define int long long
#define maxn 200005
int n;
vector<int>to[maxn];
int siz[maxn];
vector<int>f[maxn][2];
int K,g[505][2];
const int inf=1e9;
#define pb push_back
inline void dfs(int x,int pre) {
	int i,j;siz[x]=1;
	f[x][0].pb(inf),f[x][1].pb(inf),f[x][0].pb(1),f[x][1].pb(2);
	for (auto y:to[x]) if (y^pre) {
		dfs(y,x);
		for (i=1;i<=min(K,siz[x]+siz[y]);i++) g[i][0]=g[i][1]=inf; 
		for (i=1;i<=min(K,siz[x]);i++) 
			for (j=1;j<=min(K-i,siz[y]);j++){
				g[i+j][0]=min(g[i+j][0],f[x][0][i]+f[y][0][j]-i*(i+1)/2-j*(j+1)/2+(i+j)*(i+j+1)/2);
				g[i+j][1]=min(g[i+j][1],f[x][1][i]+f[y][1][j]-i*(i+1)-j*(j+1)+(i+j)*(i+j+1));
			}
		int ans1=inf,ans2=inf;
		for (j=1;j<=min(K,siz[y]);j++) ans1=min(ans1,f[y][0][j]),ans2=min(ans2,f[y][1][j]);
		for (i=1;i<=min(K,siz[x]);i++) 
			g[i][0]=min(g[i][0],f[x][0][i]+ans2),g[i][1]=min(g[i][1],f[x][1][i]+ans1);
		for (i=siz[x]+1;i<=min(K,siz[x]+siz[y]);i++) f[x][0].pb(inf),f[x][1].pb(inf);
		siz[x]+=siz[y];
		for (i=1;i<=min(K,siz[x]);i++) f[x][0][i]=g[i][0],f[x][1][i]=g[i][1]; 
		f[y][0].clear();f[y][0].shrink_to_fit();
		f[y][1].clear(),f[y][1].shrink_to_fit();
	}
}
inline void solve(void) {
	int i,x,y;
	read(n);K=min(500ll,2+(int)sqrt(n)*2);
	for (i=1;i<=n;i++) {
		to[i].clear();
		f[i][0].clear(),f[i][0].shrink_to_fit();
		f[i][1].clear(),f[i][1].shrink_to_fit();
	}
	for (i=1;i<n;i++) {
		read(x),read(y);
		to[x].push_back(y);
		to[y].push_back(x);
	}
	dfs(1,0);
	int ans=inf;
	for (i=1;i<=min(K,siz[1]);i++) ans=min(ans,min(f[1][0][i],f[1][1][i]));
	printf("%lld\n",n*(n+1)-ans);
}
signed main(void){
	int T;
	read(T);
	while (T--) solve();
	return 0;
}