rzaytsev

Problem D

Aug 21st, 2020
426
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. #include <iostream>
  2. #include <algorithm>
  3. #include <vector>
  4. #include <utility>
  5. #include <map>
  6. using namespace std;
  7. typedef vector<int> vi;
  8. typedef pair<int, int> pii;
  9. typedef long long ll;
  10. vector<bool> used;
  11. vector<vi> T;
  12. map<pii, ll> w;
  13. const ll MOD = 1'000'000'007;
  14. void dfs(int u) {
  15.    int par = -1;
  16.    ll weight = 0;
  17.    for (int v : T[u]) {
  18.        if (!used[v]) {
  19.            used[v] = true;
  20.            dfs(v);
  21.            weight += w[{u, v}];
  22.        } else { par = v; }
  23.    }
  24.    if (par != -1) {
  25.        w[{par, u}] = weight + 1;
  26.        w[{u, par}] = T.size() - weight - 1;
  27.    }
  28. }
  29. vector<ll> get_weights() {
  30.    vector<ll> weights;
  31.    for (int u = 0; u < T.size(); ++u) {
  32.        for (int v : T[u]) {
  33.            if (u < v) weights.push_back(w[{u, v}] * w[{v, u}]);
  34.        }
  35.    }
  36.    sort(weights.rbegin(), weights.rend());
  37.    return weights;
  38. }
  39. int main() {
  40.    int t;
  41.    cin >> t;
  42.    while (t--) {
  43.        int n;
  44.        cin >> n;
  45.        T = vector<vi>(n);
  46.        used = vector<bool> (n, false);
  47.        w = map<pii, ll>();
  48.        for (int i = 0; i < n - 1; ++i) {
  49.            int u, v;
  50.            cin >> u >> v;
  51.            T[u-1].push_back(v-1);
  52.            T[v-1].push_back(u-1);
  53.        }
  54.        int m;
  55.        cin >> m;
  56.        vi factors(m);
  57.        for (int& k : factors) cin >> k;
  58.        sort(factors.rbegin(), factors.rend());
  59.        dfs(0);
  60.        vector<ll> W = get_weights();
  61.        vector<ll> lab(W.size(), 1);
  62.        int f = max(0, (int) factors.size() - (int) W.size());
  63.        for (int i = 0; i < f; ++i) {
  64.            lab[0] = (lab[0] * factors[i]) % MOD;
  65.        }
  66.        for (int i = f; i < factors.size(); ++i) {
  67.            lab[i - f] = (lab[i - f] * factors[i]) % MOD;
  68.        }
  69.        ll total = 0;
  70.        for (int i = 0; i < W.size(); ++i)
  71.            total = (total + lab[i] * W[i]) % MOD;
  72.        cout << total << endl;
  73.    }
  74.    return 0;
  75. }
RAW Paste Data