oj.uz

Submission #941741

# Submission time Handle Problem Language Result Execution time Memory
941741 2024-03-09T11:37:14 Z vjudge1 Chase (CEOI17_chase) C++17
40 / 100
 419 ms 524288 KB 
#include <bits/stdc++.h>
using namespace std;

using ll = long long;
#define int ll

#define endl '\n'
#define pb push_back
using pi = pair<int, int>;
template<class T> bool ckmin(T& a, const T& b) { return b < a ? a = b, 1 : 0; }
template<class T> bool ckmax(T& a, const T& b) { return a < b ? a = b, 1 : 0; }

/*
dp[i][j] - subarborele lui i daca mai am j de aruncat
*/

const int INF = 1e16;

int32_t main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  
  int n, k;
  cin >> n >> k;
  
  vector<int> a(n + 1);
  for (int i = 1; i <= n; ++i) cin >> a[i];
  
  vector<vector<int>> adj(n + 1);
  for (int i = 1; i < n; ++i) {
    int u, v;
    cin >> u >> v;
    adj[u].pb(v);
    adj[v].pb(u);
  }
  
  vector<int> sum(n + 1);
  vector<vector<vector<int>>> max1(n + 1, vector<vector<int>>(k + 1, vector<int>(2, -INF))), max2(n + 1, vector<vector<int>>(k + 1, vector<int>(2, -INF)));
  vector<vector<int>> dp(n + 1, vector<int>(k + 1));
  function<void(int, int)> dfs = [&](int u, int p) {
    sum[u] = 0;
    for (int v : adj[u]) {
      if (v == p) continue;
      
      dfs(v, u);
      sum[u] += a[v];
    }
    
    for (int v : adj[u]) {
      if (v == p) continue;
      
      for (int i = 0; i <= k; ++i) {
        // nu folosesc
        if (dp[v][i] > max1[u][i][0]) {
          swap(max1[u][i][0], max2[u][i][0]);
          max1[u][i][0] = dp[v][i];
        } else max2[u][i][0] = max(max2[u][i][0], dp[v][i]);
        
        // folosesc
        if (i == 0) continue;
        if (dp[v][i - 1] > max1[u][i][1]) {
          swap(max1[u][i][1], max2[u][i][1]);
          max1[u][i][1] = dp[v][i - 1];
        } else max2[u][i][1] = max(max2[u][i][1], dp[v][i - 1]);
      }
    }
    
    for (int i = 0; i <= k; ++i) dp[u][i] = max({dp[u][i], max1[u][i][0], max1[u][i][1] + sum[u]});
  };
  dfs(1, 0);
  //for (int i = 0; i <= k; ++i) cout << max1[4][i][0] << " " << max2[4][i][0] << " | " << max1[4][i][1] << " " << max2[4][i][1] << endl;
  
  int ans = 0;
  function<void(int, int)> reroot = [&](int u, int p) {
    auto dpu = dp[u], dpp = dp[p];
    auto sumu = sum[u], sump = sum[p];
    auto max1u = max1[u], max2u = max2[u];
    if (p != 0) {
      // pentru p
      sum[p] -= a[u];
      for (int i = 0; i <= k; ++i) {
        // nu folosesc
        dp[p][i] = 0;
        int maxi = max1[p][i][0];
        if (dp[u][i] == maxi) maxi = max2[p][i][0];
        dp[p][i] = max(dp[p][i], maxi);
        
        // folosesc
        if (i == 0) continue;
        maxi = max1[p][i][1];
        if (dp[u][i - 1] == maxi) maxi = max2[p][i][1];
        dp[p][i] = max(dp[p][i], maxi + sum[p]);
      }
      
      // pentru u
      sum[u] += a[p];
      for (int i = 0; i <= k; ++i) {
        // nu folosesc
        dp[u][i] = 0;
        if (dp[p][i] > max1[u][i][0]) {
          swap(max1[u][i][0], max2[u][i][0]);
          max1[u][i][0] = dp[p][i];
        } else max2[u][i][0] = max(max2[u][i][0], dp[p][i]);
        dp[u][i] = max(dp[u][i], max1[u][i][0]);
        
        // folosesc
        if (i == 0) continue;
        if (dp[p][i - 1] > max1[u][i][1]) {
          swap(max1[u][i][1], max2[u][i][1]);
          max1[u][i][1] = dp[p][i - 1];
        } else max2[u][i][1] = max(max2[u][i][1], dp[p][i - 1]);
        dp[u][i] = max(dp[u][i], max1[u][i][1] + sum[u]);
      }
    }
    
    ans = max(ans, dp[u][k]);
    
    for (int v : adj[u]) {
      if (v == p) continue;
      reroot(v, u);
    }
    
    dp[u] = dpu;
    dp[p] = dpp;
    max1[u] = max1u;
    max2[u] = max2u;
    sum[u] = sumu;
    sum[p] = sump;
  };
  reroot(1, 0);
  
  cout << ans;
}

/*
good:
0 0 | -10000000000000000 -10000000000000000
16 0 | 0 0
27 0 | 16 0

bad:
0 0 | -10000000000000000 -10000000000000000
16 0 | 0 0
24 0 | 16 0

*/

Subtask #1
20.0 / 20.0

# Verdict Execution time Memory Grader output
1 Correct 0 ms 348 KB Output is correct
2 Correct 0 ms 348 KB Output is correct
3 Correct 0 ms 348 KB Output is correct
4 Correct 0 ms 348 KB Output is correct
5 Correct 0 ms 348 KB Output is correct
6 Correct 1 ms 348 KB Output is correct

Subtask #2
20.0 / 20.0

# Verdict Execution time Memory Grader output
1 Correct 0 ms 348 KB Output is correct
2 Correct 0 ms 348 KB Output is correct
3 Correct 0 ms 348 KB Output is correct
4 Correct 0 ms 348 KB Output is correct
5 Correct 0 ms 348 KB Output is correct
6 Correct 1 ms 348 KB Output is correct
7 Correct 26 ms 18780 KB Output is correct
8 Correct 2 ms 1880 KB Output is correct
9 Correct 2 ms 1116 KB Output is correct
10 Correct 20 ms 12632 KB Output is correct
11 Correct 6 ms 4188 KB Output is correct
12 Correct 2 ms 1628 KB Output is correct

Subtask #3
0 / 30.0

# Verdict Execution time Memory Grader output
1 Runtime error 419 ms 524288 KB Execution killed with signal 9
2 Halted 0 ms 0 KB -

Subtask #4
0 / 30.0

# Verdict Execution time Memory Grader output
1 Correct 0 ms 348 KB Output is correct
2 Correct 0 ms 348 KB Output is correct
3 Correct 0 ms 348 KB Output is correct
4 Correct 0 ms 348 KB Output is correct
5 Correct 0 ms 348 KB Output is correct
6 Correct 1 ms 348 KB Output is correct
7 Correct 26 ms 18780 KB Output is correct
8 Correct 2 ms 1880 KB Output is correct
9 Correct 2 ms 1116 KB Output is correct
10 Correct 20 ms 12632 KB Output is correct
11 Correct 6 ms 4188 KB Output is correct
12 Correct 2 ms 1628 KB Output is correct
13 Runtime error 419 ms 524288 KB Execution killed with signal 9
14 Halted 0 ms 0 KB -