1830D - Mex Tree - CodeForces Solution

#include <bits/stdc++.h>
using namespace std;
#define N 200005
#define ll long long
#define pb push_back
int T,n,B,dfn[N],sz[N],z[N][2],mn[N][2],dp[N][2];vector<int> e[N];
void W(int &x,int y) {x=min(x,y);}
void dfs(int u,int f)
{
	int nw=1;dfn[u]=++dfn[0];sz[u]=1;
	for(auto v:e[u]) if(v!=f) dfs(v,u),sz[u]+=sz[v];
	W(sz[u],B);for(int i=1;i<=sz[u];++i) z[i][0]=z[i][1]=1e9;
	z[1][0]=z[1][1]=0;mn[u][0]=mn[u][1]=1e9;
	for(auto v:e[u]) if(v!=f) for(int j=nw;j;--j)
	{
		for(int k=1;k<=sz[v] && j+k<=B;++k)
		{
			W(z[j+k][0],z[j][0]+dp[dfn[v]+k-1][0]);
			W(z[j+k][1],z[j][1]+dp[dfn[v]+k-1][1]);
		}z[j][0]+=mn[v][1];z[j][1]+=mn[v][0];nw=min(nw+sz[v],B);
	}
	for(int i=1;i<=sz[u];++i)
	{
		dp[dfn[u]+i-1][0]=z[i][0];dp[dfn[u]+i-1][1]=z[i][1];
		W(mn[u][0],z[i][0]+i*(i+1)/2);W(mn[u][1],z[i][1]+i*(i+1));
	}
}
void slv()
{
	scanf("%d",&n);B=500;for(int i=1;i<=n;++i) e[i].clear();
	for(int i=1,u,v;i<n;++i)
		scanf("%d %d",&u,&v),e[u].pb(v),e[v].pb(u);dfs(1,0);
	printf("%lld\n",1ll*n*(n+1)-min(mn[1][0],mn[1][1]));
}
int main() {scanf("%d",&T);while(T--) slv();return 0;}