Paths on the Tree

思路

我们从 \(|c_u-c_v|\le 1\) 切入。设节点 \(u\) 有 \(k\) 个儿子。则根据小学就学过的鸽巢原理可以得到 \(u\) 的每个儿子至少会分到 \(\left\lfloor \dfrac{c_u}{k}\right\rfloor\)条路径，最多分到 \(\left\lceil\dfrac{c_u}{k}\right\rceil\)。而且分到 \(\left\lceil\dfrac{c_u}{k}\right\rceil\) 条路径的儿子个数 \(rest=c_u\bmod k\)。由于 \(u\) 的父亲也可能给 \(u\) 多分配一条路径，所以 \(u\) 的儿子分配到的路径条数有 4 种情况：$ \left\lfloor\dfrac{c_u}{k}\right\rfloor,\left\lceil\dfrac{c_u}{k}\right\rceil,\left\lfloor\dfrac{c_u+1}{k}\right\rfloor,\left\lceil\dfrac{c_u+1}{k}\right\rceil$。其中，最小的 \(\left\lfloor\dfrac{c_u}{k}\right\rfloor\) 和最大的 \(\left\lceil\dfrac{c_u+1}{k}\right\rceil\) 相差不会超过 1，所以一个节点分配到的节点数只有两种情况。

然后我们来决定这 \(rest\) 个多分配一条路径的机会分给那些儿子。显然哪些儿子多分配后增加的收益最大就分配给谁，用递归去求解。

Code

#include<bits/stdc++.h>
#define int long long
#define push_back emplace_back
using namespace std;
const int N = 2e5+25;
const int INF = 1e9+7;
int T;
int n,k,s[N];
vector<int>g[N];//存边
map<int,int>f[N];//记录每个节点带来的收益
int dfs(int u,int m){
	if(m==0)return 0;
	if(f[u].count(m))return f[u][m];
	if(g[u].size()==0)return f[u][m]=m*s[u];
	int t=m/g[u].size(),rest=m%g[u].size();
	int res=0;
	for(auto v:g[u])res+=dfs(v,t);
	if(rest){
		priority_queue<int>q;//将增加的收益用优先队列来维护
		for(auto v:g[u])q.emplace(dfs(v,t+1)-dfs(v,t));
		for(int i=0;i<rest;i++)res+=q.top(),q.pop();
	}
	res+=m*s[u];
	return f[u][m]=res;
}
void solve(){
	cin>>n>>k;
	for(int i=1;i<=n;i++)
		g[i].clear(),f[i].clear();
	for(int i=2,u;i<=n;i++)
		cin>>u,g[u].push_back(i);
	for(int i=1;i<=n;i++)cin>>s[i];
	cout<<dfs(1,k)<<'\n';
}
signed main(){
	ios::sync_with_stdio(false);
	cin.tie(0),cout.tie(0);
	cin>>T;
	while(T--)solve();
	return 0;
}