【題解】Codeforces 1401D. Maximum Distributed Tree

【題目敘述】http://codeforces.com/contest/1401/problem/D

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
 
int t, n, m, a, b, mod = 1e9+7;
long long p[100005], sz[100005];
vector <int> v[100005];
vector <long long> w;
 
void dfs(int x, int pre){
    sz[x] = 0;
    for (int i:v[x]){
        if (i == pre) continue;
        dfs(i, x);
        w.push_back(sz[i]*(n-sz[i]));
        sz[x] += sz[i];
    }
    sz[x]++;
    return;
}
 
int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cin >> t;
    while (t--){
        cin >> n;
        for (int i = 1; i <= n; i++){
            v[i].clear();
        }
        for (int i = 1; i < n; i++){
            cin >> a >> b;
            v[a].push_back(b);
            v[b].push_back(a);
        }
        w.clear();
        dfs(1, 0);
        sort(w.begin(), w.end());
        cin >> m;
        for (int i = 0; i < m; i++){
            cin >> p[i];
        }
        if (m <= n-1){
            for (int i = m; i < n-1; i++){
                p[i] = 1;
            }
            sort(p, p+n-1);
        }
        else{
            sort(p, p+m);
            for (int i = n-1; i < m; i++){
                p[n-2] *= p[i];
                p[n-2] %= mod;
            }
        }
        long long ans = 0;
        for (int i = 0; i < n-1; i++){
            ans += w[i]*p[i];
            ans %= mod;
        }
        cout << ans << "\n";
    }
}
分享本文 Share with friends