Codeforces Round #665 (Div. 2) D - Maximum Distributed Tree（dfs记录子节点+贪心）

Codeforces Round #665 (Div. 2) D - Maximum Distributed Tree（dfs记录子节点+贪心）

题目

http://codeforces.com/contest/1401/problem/D

题意

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

题解

先记录出每个边被计算出来的次数，siz[i] = son[i]*(n-son[i])
再给m个质数排序，把n-1个数后面的数合并，然后一个个乘以每个边被计算的次数，记得取mod。

代码

#include<iostream>
#include<stdio.h>
#include<string.h>
#include<math.h>
#include<algorithm>
#include<vector>
#include<map>
#include<set>
#include<queue>
#include<stack>
using namespace std;
#define ll long long
#define ld long double
#define INF 0x3f3f3f3f
#define mem(a,b) memset(a,b,sizeof(a))
#define PI acos(-1)
const int maxn = 1e5+5;
const ll mod = 1e9+7;
const double eps = 1e-6L;
vector<int> g[maxn];
int son[maxn];
ll k[maxn],siz[maxn];
void dfs(int u, int fa)
{
	son[u]=1;
	if(g[u].size()==1&&fa!=-1)
	{
		return;
	}
	for(auto x:g[u])
	{
		if(x!=fa)
		{
			dfs(x,u);
			son[u]+=son[x];
		}
	}
}
int main()
{
	int T=1;
	scanf("%d",&T);
	while(T--)
	{
		int n,m;
		scanf("%d",&n);
		for(int i=1;i<=n;i++)	g[i].clear();
		for(int i=0;i<n-1;i++)
		{
			int u,v;
			scanf("%d %d",&u,&v);
			g[u].push_back(v);
			g[v].push_back(u);
		}
		scanf("%d",&m);
		for(int i=0;i<m;i++)
			scanf("%d",&k[i]);
		sort(k,k+m);
		dfs(1,-1);
		for(int i=1;i<n;i++)
		{
			siz[i] = 1ll*son[i+1]*(n-son[i+1]);
		}
		sort(siz+1,siz+n);
		ll ans=0;
		while(m>n-1)
		{
			k[m-2] = k[m-2]*k[m-1]%mod;
			m--;
		 } 
		for(int i=n-1,j=m-1;i>=1;i--,j--)
		{
			if(j>=0)	ans+=siz[i]*k[j]%mod;
			else	ans+=siz[i];
			ans %= mod;
		}
		printf("%lld\n",ans);
	}
	return 0;
}