CF1401D-Maximum Distributed Tree

题目

题意

给出一棵树和一个数k，请你给这棵树上的每条边分配一个边权，使得这棵树上的每两个节点之间的路径总和最大，同时边权中1的数量尽可能少，所有边权的乘积等于k。

分析

计算贡献，即每条边选中的次数。将边按照贡献从小到大排序，优先将大的质因子和贡献大的边进行组合。
用dfs跑一遍算出每条边后面有多少个节点，然后贡献=该边前面节点数*该边后面节点数。
分两种情况，质数不够用1来凑；质数多了，把最后几个大的质数合并为一个数，并分配给贡献最多的边。

代码

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const ll mod = 1e9 + 7;
const int N = 1e5 + 10;

vector<int>vt[N];
int p[N], a[N];
ll edge[N];//存每条边的贡献

//dfs算每条边后面有多少个节点，即a数组
void dfs(int now, int pre) {
	a[now] = 1;
	for (int i = 0; i < vt[now].size(); i++) {
		int v = vt[now][i];
		if (v == pre)
			continue;
		dfs(v, now);
		a[now] += a[v];
		//cout << "a[now]:" << now << ' ' << a[now] << endl;
	}
}

int main() {
	int T;
	cin >> T;
	while (T--) {
		int n, u, v, m;
		cin >> n;
		for (int i = 0; i <= n; i++) {
			vt[i].clear();
			a[i] = 0;
			p[i] = 0;
			edge[i] = 0;
		}
		for (int i = 0; i < n - 1; i++) {
			cin >> u >> v;
			vt[u].push_back(v);
			vt[v].push_back(u);
		}
		cin >> m;
		for (int i = 1; i <= m; i++)
			cin >> p[i];
		sort(p + 1, p + m + 1);
		dfs(1, -1);
		ll ans = 0;
		/*for (int i = 1; i <= n; i++)
			cout << a[i] << endl;*/
		for (int i = 1; i <= n; i++) {
			//贡献=该边前面节点数*该边后面节点数
			edge[i - 1] = ll(a[i]) * (n - a[i]);
			//cout << "edge[i-1]" << i - 1 << ' ' << edge[i - 1] << endl;
		}
		sort(edge + 1, edge + n);
		//分情况计算ans
		if (n - 1 >= m) {
			int t = n - 1 - m;
			for (int i = 1; i <= t; i++)
				ans = (ans + edge[i] % mod) % mod;
			for (int i = t + 1; i < n; i++)
				ans = (ans + edge[i] * p[i - t] % mod) % mod;
		}
		else {
			ll t = 1;
			for (int i = n - 1; i <= m; i++)
				t = t * p[i] % mod;
			for (int i = 1; i < n - 1; i++) {
				ans = (ans + edge[i] * p[i] % mod) % mod;
			}
			ans = (ans + edge[n-1] * t % mod) % mod;
		}
		cout << ans << endl;
	}
	return 0;
}