自动搬运

查看原文

来自洛谷，原作者为

	RAYMOND_7 没有卓子张@Reunite是万万不能的

搬运于2025-08-24 22:11:44，当前版本为作者最后更新于2025-01-25 22:25:16，作者可能在搬运后再次修改，您可在原文处查看最新版

自动搬运只会搬运当前题目点赞数最高的题解，您可前往洛谷题解查看更多

以下是正文

首先发现一个结论：最优解存在一个点作为根，其每个子树要不然是内向树，要不然是外向树。

调整法可以证明，每个子树内的方向一定是统一的。接下来就很好办了，到根的贡献是非常好算的，每个子树内部的贡献也是确定的，现在的问题转化为将子树划分成两个集合使得两个集合求和后的乘积最大。

又基本不等式可知一定是接近总和一半更优，由题意 $\forall deg_u \le 36$ 想到折半搜索，到这里我们已经获得 $O(n\times 2^\frac{D}{2})$ 的做法，直接实现可以得到 $80$ 分。

发现一个有趣的性质，就是除了带权中心外一定存在一个子树满足权值达到一半，那么一定是这个子树单独划分为一个集合，否则更劣，于是非带权中心不需要折半搜索，从而这部分复杂度锐减至 $O(2^\frac{D}{2}+n)$。

还有一个问题就是子树内的贡献怎么算，每次换根应该是可以直接维护的，但是我比较懒写了对边的两个方向进行记忆化搜索，这个复杂度本身不对可以被菊花图卡，但是这个题度数很小所以最坏是 $O(nD\log n)/O(nD)$ 的，$O(\log n)$ 是 map 存状态的复杂度。

代码很简单：

#include <cstdio>
#include <algorithm>
#include <cstring>
#include <vector>
#include <map>
#include <cassert>
#include <string>
#include <iostream>
#include <queue>

#define int long long
#define PII pair<int,int>
#define x first
#define y second
#define For(i,l,r) for(int i=l;i<=r;i++)
#define Rep(i,l,r) for(int i=l;i>=r;i--)

using namespace std;
bool _A_A;

void in(int &x)
{
	char c=getchar();int op=1;
	while(c>'9'||c<'0')op*=(c=='-'?-1:1),c=getchar();
	for(x=0;c>='0'&&c<='9';c=getchar())x=x*10+c-'0';x*=op;
}

const int N=2e5+50;

int n,m,a[N],p[N],sz[N],f[N],son[N],num;
vector<int>G[N];
int o[N],cnt;
map<PII,int> mp;

void dfs(int u,int pa)
{
	sz[u]=0;f[u]=0;
	for(int v:G[u])if(v!=pa)
	{
		dfs(v,u);
		sz[u]+=sz[v];f[u]+=f[v];
	}
	f[u]+=sz[u]*a[u];
	sz[u]+=a[u];
}

int gsz(int u,int pa)
{
	if(p[u]==pa)return sz[u];return sz[1]-sz[pa];
}

int work(int u,int pa)
{
	if(mp.count({u,pa}))return mp[{u,pa}];
	int s=0;
	for(int v:G[u])if(v!=pa)
	{
		s+=work(v,u)+a[u]*gsz(v,u);
	}
	return mp[{u,pa}]=s;
}

bool _B_B;
signed main()
{
	in(n);For(i,1,n)in(a[i]);For(i,2,n)in(p[i]),G[p[i]].push_back(i),G[i].push_back(p[i]);
	int ans=0;dfs(1,0);
	For(u,1,n)
	{
		cnt=0;int res=0,as=0;
		for(int v:G[u])o[++cnt]=gsz(v,u),res+=work(v,u)+gsz(v,u)*a[u];
		num=0;int all=sz[1]-a[u];
		int dec=0;
		For(i,1,cnt)dec=max(dec,o[i]);
		if(dec<=all/2)
		{
			auto dfs1=[&](auto dfs1,int k,int s)
			{
				if(k>cnt/2)
				{
					son[++num]=s;return ;
				}
				dfs1(dfs1,k+1,s+o[k]);
				dfs1(dfs1,k+1,s);
			};
			auto dfs2=[&](auto dfs2,int k,int s)
			{
				if(k<=cnt/2)
				{
					int t=lower_bound(son+1,son+num+1,all/2-s)-son;
					if(t<=num)as=max(as,(son[t]+s)*(all-son[t]-s));
					return ;
				}
				dfs2(dfs2,k-1,s+o[k]);
				dfs2(dfs2,k-1,s);
			};
			dfs1(dfs1,1,0);sort(son+1,son+num+1);
			dfs2(dfs2,cnt,0);
		}
		else as=dec*(all-dec);
		res+=as;ans=max(ans,res);
	}
	For(i,1,n)ans+=a[i]*(a[i]-1);cout<<ans<<"\n";
	cerr<<(&_B_B-&_A_A)*1.0/1024/1024<<endl;
	return 0;
}