Codeforces D. Maximum Distributed Tree (树dfs & 贪心)(Round #665 Div.2)

传送门

题意： 给出一棵n个点的树，想让给每条边赋值，使得所有边权的乘积为k。因为k非常大，所有将k拆分成m个因子，换而言之就是n-1条边权的乘积得等于m个因子的乘积。
定义f(x, y)为点x到点y的简单路径所经过的所有边的权值和，现问怎么安排边权，才能让所有简单路径的 f 值和 (即下图所示公式)最大，并得到该max值。


思路：

• 比赛的时候相叉了，后面快结束了才行到 idea， 但对树的相关操作不熟就这么错过了上高分分的机会，呜呜太菜了~
• 若m < n-1，那么剩下的边用1来填充。
• dfs找到每条边的经过次数，然后贪心一下将大的因子赋值给经过次数最多的边贡献就最多。
• 至于经过次数的计算，可以将 u - v 的边去掉，就会形成两个连通块，该边的经过次数就为两个连通块大小的乘积 size(u) * (n- size(v))。

代码实现：

#include
#define endl '\n'
#define null NULL
#define ll long long
#define int long long
#define pii pair
#define lowbit(x) (x &(-x))
#define ls(x) x<<1
#define rs(x) (x<<1+1)
#define me(ar) memset(ar, 0, sizeof ar)
#define mem(ar,num) memset(ar, num, sizeof ar)
#define rp(i, n) for(int i = 0, i < n; i ++)
#define rep(i, a, n) for(int i = a; i <= n; i ++)
#define pre(i, n, a) for(int i = n; i >= a; i --)
#define IOS ios::sync_with_stdio(0); cin.tie(0);cout.tie(0);
const int way[4][2] = {{1, 0}, {-1, 0}, {0, 1}, {0, -1}};
using namespace std;
const int  inf = 0x7fffffff;
const double PI = acos(-1.0);
const double eps = 1e-6;
const ll   mod = 1e9 + 7;
const int  N = 5e5 + 5;

int t, n, m;
int s[N], c[N], p[N];
int h[N], tt;

struct node{
	int to, ne;
}g[N];
void add(int u,int v){
	g[++ tt]=(node){v, h[u]};
	h[u] = tt;
}

void dfs(int u,int fa){
	s[u] = 1;
	for(int i = h[u]; i; i = g[i].ne){
		int v = g[i].to;
		if(v == fa)	continue;
		dfs(v, u);
		s[u] += s[v];
		c[i/2] = s[v]*(n-s[v]);
	}
}

signed main()
{
    IOS;
	cin >> t;
	while(t --){
		cin >> n; tt = 1;
		for(int i = 0; i <= n*2;i ++) h[i] = 0, s[i] = 0;
		for(int i = 1; i <= m; i ++) p[i] = 0;
		for(int i = 1; i < n; i ++){
			int u, v; cin >> u >> v;
			add(u, v); add(v, u);
		}
		dfs(1,0);
		cin >> m;
		for(int i = 1; i <= m; i ++) cin >> p[i];
		sort(p+1, p+m+1); sort(c+1,c+n);
		int ans = 0;
		if(m <= n-1){
			for(int i = n-1, j = m; j >= 1; i --, j --)
				ans = (ans + c[i]*p[j]%mod )%mod;
			for(int i = n-m-1; i >= 1; i --)
				ans = (ans + c[i])%mod;

		}
		else{
			for(int i = m-1; i >= n-1;i --)
				p[i] = p[i]*p[i+1]%mod;
			for(int i = n-1; i >= 1; i --)
				ans = (ans + c[i]*p[i]%mod )%mod;
		}
		cout << ans << endl;
	}
	return 0;
}