Codeforces1401 D. Maximum Distributed Tree（DFS）

本文主要是介绍Codeforces1401 D. Maximum Distributed Tree（DFS），希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

You are given a tree that consists of 𝑛 nodes. You should label each of its 𝑛−1 edges with an integer in such way that satisfies the following conditions:

each integer must be greater than 0;
the product of all 𝑛−1 numbers should be equal to 𝑘;
the number of 1-s among all 𝑛−1 integers must be minimum possible.
Let’s define 𝑓(𝑢,𝑣) as the sum of the numbers on the simple path from node 𝑢 to node 𝑣. Also, let ∑𝑖=1𝑛−1∑𝑗=𝑖+1𝑛𝑓(𝑖,𝑗) be a distribution index of the tree.

Find the maximum possible distribution index you can get. Since answer can be too large, print it modulo 109+7.

In this problem, since the number 𝑘 can be large, the result of the prime factorization of 𝑘 is given instead.

Input
The first line contains one integer 𝑡 (1≤𝑡≤100) — the number of test cases.

The first line of each test case contains a single integer 𝑛 (2≤𝑛≤105) — the number of nodes in the tree.

Each of the next 𝑛−1 lines describes an edge: the 𝑖-th line contains two integers 𝑢𝑖 and 𝑣𝑖 (1≤𝑢𝑖,𝑣𝑖≤𝑛; 𝑢𝑖≠𝑣𝑖) — indices of vertices connected by the 𝑖-th edge.

Next line contains a single integer 𝑚 (1≤𝑚≤6⋅104) — the number of prime factors of 𝑘.

Next line contains 𝑚 prime numbers 𝑝1,𝑝2,…,𝑝𝑚 (2≤𝑝𝑖<6⋅104) such that 𝑘=𝑝1⋅𝑝2⋅…⋅𝑝𝑚.

It is guaranteed that the sum of 𝑛 over all test cases doesn’t exceed 105, the sum of 𝑚 over all test cases doesn’t exceed 6⋅104, and the given edges for each test cases form a tree.

Output
Print the maximum distribution index you can get. Since answer can be too large, print it modulo 109+7.

Example
inputCopy
3
4
1 2
2 3
3 4
2
2 2
4
3 4
1 3
3 2
2
3 2
7
6 1
2 3
4 6
7 3
5 1
3 6
4
7 5 13 3
outputCopy
17
18
286
Note
In the first test case, one of the optimal ways is on the following image:

In this case, 𝑓(1,2)=1, 𝑓(1,3)=3, 𝑓(1,4)=5, 𝑓(2,3)=2, 𝑓(2,4)=4, 𝑓(3,4)=2, so the sum of these 6 numbers is 17.

In the second test case, one of the optimal ways is on the following image:

In this case, 𝑓(1,2)=3, 𝑓(1,3)=1, 𝑓(1,4)=4, 𝑓(2,3)=2, 𝑓(2,4)=5, 𝑓(3,4)=3, so the sum of these 6 numbers is 18.

思路：
只要算出每个边的贡献，再把质因数排个序依次匹配即可。
如果质因数小于n-1，那么添加1。
如果大于n-1，那么合并最后的几个质因数。

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <vector>
#include <queue>using namespace std;typedef long long ll;
const int maxn = 4e5 + 7;
const int mod = 1e9 + 7;
const int INF = 0x3f3f3f3f;int head[maxn],nex[maxn],to[maxn],tot;
ll dp[maxn],a[maxn];
ll siz[maxn];
int n;void add(int x,int y) {to[++tot] = y;nex[tot] = head[x];head[x] = tot;
}void init() {tot = 0;for(int i = 1;i <= n;i++) {head[i] = dp[i] = 0;}
}void dfs(int x,int fa) {siz[x] = 1;for(int i = head[x];i;i = nex[i]) {int v = to[i];if(v == fa) continue;dfs(v,x);siz[x] += siz[v];dp[(i + 1) / 2] = 1ll * siz[v] * (n - siz[v]);}
}int main() {int T;scanf("%d",&T);while(T--) {scanf("%d",&n);init();for(int i = 1;i < n;i++) {int x,y;scanf("%d%d",&x,&y);add(x,y);add(y,x);}int m;scanf("%d",&m);for(int i = 1;i <= m;i++) {scanf("%lld",&a[i]);}dfs(1,-1);while(m < n - 1) {a[++m] = 1;}sort(a + 1,a + 1 + m);sort(dp + 1,dp + 1 + n - 1);if(m > n - 1) {for(int j = n;j <= m;j++) {a[n - 1] = a[n - 1] * a[j] % mod;}}ll ans = 0;for(int i = 1;i <= n - 1;i++) {ans = (ans + dp[i] * a[i] % mod) % mod;}printf("%lld\n",ans);}return 0;
}