Advertisement
Guest User

Untitled

a guest
Oct 22nd, 2019
122
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 3.59 KB | None | 0 0
  1. #include <bits/stdc++.h>
  2.  
  3. using namespace std;
  4.  
  5. #define forn(a, e) for (int a = 0; a < (int)(e); a++)
  6. #define forr(a, s, e) for (int a = s; a < (int)(e); a++)
  7. #define fore(e, a) for (auto& e : a)
  8.  
  9. #ifdef LOCAL
  10. #define logv(a) {cerr << #a << " = "; fore(e, a) {cerr << e << " ";} cerr << "\n";}
  11. #define logvp(a) {cerr << #a << " = "; fore(e, a) {cerr << "(" << e.first << ", " << e.second << ") ";} cerr << "\n";}
  12. #define logvv(a) {cerr << #a << " = \n"; fore(r, a) { fore(e, r) {cerr << e << " ";} cerr << "\n";} }
  13. #define logvf(a, field) {cerr << #a"."#field << " = \n"; fore(e, a) { cerr << e.field << " ";} cerr << "\n"; }
  14. #define logs(a) cerr << #a << " = " << (a) << "\n";
  15. #define logss(a, b) cerr << #a << " = " << (a) << ", " << #b << " = " << (b) << "\n";
  16. #define logp(a) cerr << #a << " = " << "(" << a.first << ", " << a.second << ")" << "\n";
  17. #define cond(pred, stmt) if (pred) { stmt }
  18. #else
  19. #define logv(a)
  20. #define logvp(a)
  21. #define logvv(a)
  22. #define logvf(a, field)
  23. #define logs(a)
  24. #define logss(a, b)
  25. #define logp(a)
  26. #define cond(pred, stmt)
  27. #endif
  28.  
  29. using i64 = long long;
  30. using iip = pair<i64, i64>;
  31. using ivec = vector<int>;
  32. using llvec = vector<i64>;
  33. using svec = vector<string>;
  34.  
  35. template<typename T, typename Dim>
  36. auto make_vec(T value, Dim dim) { return vector<T>(dim, value); }
  37. template<typename T, typename Dim1, typename... Dim>
  38. auto make_vec(T value, Dim1 dim1, Dim... dims) { return make_vec(make_vec(value, dims...), dim1); }
  39.  
  40. template<typename T>
  41. bool uax(T& v, const T& newv) { if (v < newv) { v = newv; return true; } else return false; }
  42. template<typename T>
  43. bool uin(T& v, const T& newv) { if (v > newv) { v = newv; return true; } else return false; }
  44.  
  45. template<typename T>
  46. istream& operator>>(istream& is, vector<T>& c) { for (auto& e : c) is >> e; return is; }
  47.  
  48. template<typename ...T>
  49. istream& read(T&... args) { return (cin >> ... >> args); }
  50.  
  51. static mt19937 rande(123123);
  52. template<typename T>
  53. T rand_int(T from, T to) { uniform_int_distribution<T> distr(from, to); return distr(rande); }
  54.  
  55. int main() {
  56. ios_base::sync_with_stdio(false);
  57. cin.tie(NULL);
  58.  
  59. int n, k;
  60. while (read(n, k)) {
  61. ivec a(n);
  62. cin >> a;
  63.  
  64. vector<vector<int>> g(n);
  65. forn(i, n - 1) {
  66. int u, v;
  67. read(u, v);
  68. u--; v--;
  69. g[u].push_back(v);
  70. g[v].push_back(u);
  71. }
  72.  
  73. auto idp = make_vec(0, n, max(n, k + 1));
  74. auto odp = idp;
  75.  
  76. function<void(int, int)> dfs = [&](int v, int p) {
  77.  
  78. fore(u, g[v]) {
  79. if (u != p) {
  80. dfs(u, v);
  81. }
  82. }
  83.  
  84. idp[v][0] = a[v];
  85. forr(i, 1, n) {
  86. fore(u, g[v]) {
  87. idp[v][i] += idp[u][i - 1];
  88. }
  89. }
  90.  
  91. if (p != -1) {
  92. forr(i, 1, n - 1) {
  93. odp[v][i] = odp[p][i + 1];
  94. }
  95. odp[v][1] = a[p];
  96. }
  97.  
  98. forn(i, n - 1) {
  99. fore(u, g[v]) {
  100. if (u == p) {
  101. continue;
  102. }
  103. odp[u][i + 1] += idp[v][i];
  104. if (i > 0) {
  105. odp[u][i + 1] -= idp[u][i - 1];
  106. }
  107. }
  108. }
  109. };
  110.  
  111. dfs(0, -1);
  112.  
  113. logvv(idp);
  114. logvv(odp);
  115. int maxs = 0;
  116.  
  117. forn(i, n) {
  118. uax(maxs, idp[i][k + 1] + odp[i][k + 1]);
  119. }
  120. cout << maxs << endl;
  121. }
  122. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement