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;
}