fork download 
  1. #include <bits/stdc++.h>
  2.  
  3. typedef long long int ll;
  4. using namespace std;
  5.  
  6. ll dfs(int s, int p, vector<ll>& w, vector<unordered_set<int>>& tree, unordered_map<int, int>& parent, map<pair<ll, ll>, int>& arr) {
  7.     parent[s] = p;
  8.     ll sum = 0;
  9.     for (auto v : tree[s]) {
  10.         if (v == p)
  11.             continue;
  12.         sum += dfs(v, s, w, tree, parent, arr);
  13.     }
  14.     sum += w[s];
  15.     arr[{sum, w[s]}] = s;
  16.     return sum;
  17. }
  18.  
  19. ll dfs_sum(int s, int p, vector<ll>& w, vector<unordered_set<int>>& tree) {
  20.     ll sum = 0;
  21.     for (auto v : tree[s]) {
  22.         if (v == p)
  23.             continue;
  24.         sum += dfs_sum(v, s, w, tree);
  25.     }
  26.     sum += w[s];
  27.     return sum;
  28. }
  29.  
  30. int main() {
  31.     int n, k;
  32.     cin >> n >> k;
  33.  
  34.     vector<ll> w(n+1);
  35.     for (int i = 1; i <= n; i++)
  36.         cin >> w[i];
  37.  
  38.     vector<unordered_set<int>> tree(n+1);
  39.     int u, v;
  40.     for (int i = 1; i <= n-1; i++) {
  41.         cin >> u >> v;
  42.         tree[u].insert(v);
  43.         tree[v].insert(u);
  44.     }
  45.  
  46.     ll last_sum = LLONG_MIN;
  47.  
  48.     for (int i = 1; i <= k; i++) {
  49.         map<pair<ll, ll>, int> arr;
  50.         unordered_map<int, int> parent;
  51.  
  52.         ll sum = dfs(1, -1, w, tree, parent, arr);
  53.         last_sum = max(last_sum, sum);
  54.         if (arr.begin()->first.first >= 0)
  55.             break;
  56.  
  57.         int s = arr.begin()->second;
  58.         if (s == 1) {
  59.             cout << 0 << "\n";
  60.             return 0;
  61.         }
  62.         tree[parent[s]].erase(s);
  63.         tree[s].erase(parent[s]);
  64.     }
  65.  
  66.     cout << max(last_sum, dfs_sum(1, -1, w, tree)) << "\n";
  67. }
Runtime error #stdin #stdout 0s 4284KB
comments (?)
stdin 
Standard input is empty
stdout 
Standard output is empty
https://ideone.com/fsNHc2
5 2
1 1 -1 -1 -1
1 2
2 3
4 1
4 5
language:
C++14 (gcc 8.3)
created:
4 years ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!