Codeforces Round #665 (Div. 2) D. Maximum Distributed Tree (dfs计数,树)

• 题意:给你含有\(n\)个节点,\(n-1\)条边的树,以及\(m\)个质数和\(1\),你需要在这\(m\)个质数和一个\(1\)选择数(质数只能选一次,\(1\)可以多选)给\(n-1\)条边赋值,求所有简单路径的边权和.
• 题解:很简单,对于每条边,我们看它左右有多少个点,右边有多少点,左边点数x右边点数就是包含这条边的简单路径数,也就是说这条边权要计算的次数,我们一定会把最大的边权赋给简单路径数最多的边,所以我们可以直接dfs求每个点的子节点个数\(son[u]\)(右边的点数),那么\(n-son[u]\)就是左边的节点个数,然后可以直接乘边权算答案.要特别注意\(m\ge n\)时,我们将\(p[n-1]\)后面的数都乘给\(p[n-1]\)即可.
• 代码;

#include <bits/stdc++.h>
#define ll long long
#define fi first
#define se second
#define pb push_back
#define me memset
#define rep(a,b,c) for(int a=b;a<=c;++a)
#define per(a,b,c) for(int a=b;a>=c;--a)
const int N = 1e6 + 10;
const int mod = 1e9 + 7;
const int INF = 0x3f3f3f3f;
using namespace std;
typedef pair<int,int> PII;
typedef pair<ll,ll> PLL;
ll gcd(ll a,ll b) {return b?gcd(b,a%b):a;}
ll lcm(ll a,ll b) {return a/gcd(a,b)*b;}

#define int long long

int t;
int n;
int u,v;
vector<int> e[N];
int m;
int p[N];
int son[N];
vector<int> res;

void dfs(int u,int fa){
	son[u]=1;
	for(int w : e[u]){
		if(w == fa) continue;
		dfs(w,u);
		son[u]+=son[w];
	}
}

void cal(){
	/*
	for(int w : e[u]){
		if(w == fa) continue;
		int cur=son[w]*(n-son[w]);
		res.pb(cur);
		cal(w,u);
	}
	*/
	rep(i,2,n){
		int cur=son[i]*(n-son[i]);
		res.pb(cur);
	}
}

signed main() {
    ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);
	cin>>t;
	while(t--){
		cin>>n;

		rep(i,1,n) e[i].clear();
		res.clear();

		rep(i,1,n-1){
			cin>>u>>v;
			e[u].pb(v);
			e[v].pb(u);
		}

		cin>>m;

		rep(i,1,m){
			cin>>p[i];
		}

		dfs(1,0);
		cal();

		sort(p+1,p+1+m);

		n--;

		if(m>n){
			rep(i,n+1,m) p[n]=(p[n]*p[i])%mod;
			m=n;
		}

		sort(res.begin(),res.end(),greater<int>());

		ll ans=0;
		rep(i,0,(int)res.size()-1){
			ans=(ans+res[i]*max(1ll,p[m--]))%mod;
		}

		cout<<ans<<'\n';

	}

    return 0;
}