fork download
copy 
  1. #include <bits/stdc++.h>
  2. #define int long long
  3. #define endl '\n'
  4. #define MASK(n) (1LL << n)
  5. #define PhTrNghia "Chase"
  6.  
  7. using namespace std;
  8.  
  9. const int maxn = 2e5 + 5;
  10. const int maxk = 205;
  11. const int inf = 1e18;
  12.  
  13. int n, k;
  14. int a[maxn];
  15. int dp[maxn][maxk][2], s[maxn];
  16. vector <int> g[maxn];
  17.  
  18. namespace sub2{
  19.     int ans = 0;
  20.  
  21.     void dfs(int u, int p){
  22.         for (int i = 0; i <= k; i++) ans = max(ans, dp[u][i][1]);
  23.  
  24.         for (int v: g[u]){
  25.             if (v == p) continue;
  26.             dp[v][0][0] = dp[u][0][0];
  27.  
  28.             for (int i = 1; i <= k; i++){
  29.                 dp[v][i][0] = max({dp[v][i][0], dp[u][i][0], dp[u][i][1]});
  30.  
  31.                 int val = s[v] - a[u];
  32.                 dp[v][i][1] = max({dp[v][i][1], dp[u][i-1][1] + val, dp[u][i-1][0] + val});
  33.             }
  34.  
  35.             dfs(v, u);
  36.         }
  37.  
  38.         for (int i = 0; i <= k; i++){
  39.             dp[u][i][0] = -inf;
  40.             dp[u][i][1] = -inf;
  41.         }
  42.     }
  43.  
  44.     void solve(){
  45.         for(int i = 1; i <= n; i++){
  46.             dp[i][0][0] = 0;
  47.             dp[i][1][1] = s[i];
  48.             dfs(i, 0);
  49.         }
  50.         cout << ans;
  51.     }
  52. }
  53.  
  54. namespace sub3{
  55.     int ans = 0;
  56.  
  57.     void dfs(int u, int p){
  58.         for (int i = 0; i <= k; i++) ans = max(ans, dp[u][i][1]);
  59.  
  60.         for (int v: g[u]){ 
  61.             if (v == p) continue;
  62.  
  63.             dp[v][0][0] = dp[u][0][0];
  64.  
  65.             for (int i = 1; i <= k; i++){
  66.                 dp[v][i][0] = max({dp[v][i][0], dp[u][i][0], dp[u][i][1]});
  67.  
  68.                 int val = s[v] - a[u];
  69.                 dp[v][i][1] = max({dp[v][i][1], dp[u][i-1][1] + val, dp[u][i-1][0] + val});
  70.             }
  71.  
  72.             dfs(v, u);
  73.         }
  74.     }
  75.  
  76.     void solve(){
  77.         memset(dp, -0x3f, sizeof dp);
  78.         dp[1][0][0] = 0;
  79.         dp[1][1][1] = s[1];
  80.  
  81.         dfs(1, 0);
  82.  
  83.         cout << ans;
  84.     }
  85. }
  86.  
  87. signed main(){
  88.     ios_base::sync_with_stdio(false);
  89.     cin.tie(0); cout.tie(0);
  90.     if (fopen(PhTrNghia".INP", "r")){
  91.         freopen(PhTrNghia".INP", "r", stdin);
  92.         freopen(PhTrNghia".OUT", "w", stdout);
  93.     }
  94.  
  95.     cin >> n >> k;
  96.  
  97.     for (int i = 1; i <= n; i++) cin >> a[i];
  98.  
  99.     for (int i = 1; i < n; i++){
  100.         int u, v; 
  101.         cin >> u >> v;
  102.         g[u].push_back(v);
  103.         g[v].push_back(u);
  104.         s[u] += a[v];
  105.         s[v] += a[u];
  106.     }
  107.  
  108.  
  109.     if (n <= 1000) sub2::solve(); else sub3::solve();
  110.  
  111.     return 0;
  112. }
Success #stdin #stdout 0.01s 8484KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
Standard output is empty
https://ideone.com/m8OAAP
language:
C++14 (gcc 8.3)
created:
1 day 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!