oj.uz

Submission #941828

# Submission time Handle Problem Language Result Execution time Memory
941828 2024-03-09T14:10:31 Z rolandpetrean Chase (CEOI17_chase) C++17
100 / 100
 424 ms 466492 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 = 1e14;
const int N = 1e5 + 5, K = 105;

int n, k;
int a[N];
vector<int> adj[N];

int sum[N];
int max1[N][K][2], max2[N][K][2];
int dp[N][K];

void 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]});
}

int ans = 0;
void 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];*/
  int sump = sum[p];
  vector<int> dpp(k + 1);
  for (int i = 0; i <= k; ++i) dpp[i] = dp[p][i];
  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);
  }
  
  sum[p] = sump;
  for (int i = 0; i <= k; ++i) dp[p][i] = dpp[i];
  /*dp[u] = dpu;
  dp[p] = dpp;
  max1[u] = max1u;
  max2[u] = max2u;
  sum[u] = sumu;
  sum[p] = sump;*/
}

int32_t main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  
  cin >> n >> k;
  
  for (int i = 1; i <= n; ++i) cin >> a[i];
  for (int i = 1; i < n; ++i) {
    int u, v;
    cin >> u >> v;
    adj[u].pb(v);
    adj[v].pb(u);
  }
  
  for (int i = 0; i < N; ++i) {
    for (int j = 0; j < K; ++j) {
      max1[i][j][0] = max1[i][j][1] = max2[i][j][0] = max2[i][j][1] = -INF;
    }
  }
  
  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;
  
  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 133 ms 334928 KB Output is correct
2 Correct 106 ms 334928 KB Output is correct
3 Correct 103 ms 334932 KB Output is correct
4 Correct 102 ms 334928 KB Output is correct
5 Correct 110 ms 334760 KB Output is correct
6 Correct 118 ms 334848 KB Output is correct

Subtask #2
20.0 / 20.0

# Verdict Execution time Memory Grader output
1 Correct 133 ms 334928 KB Output is correct
2 Correct 106 ms 334928 KB Output is correct
3 Correct 103 ms 334932 KB Output is correct
4 Correct 102 ms 334928 KB Output is correct
5 Correct 110 ms 334760 KB Output is correct
6 Correct 118 ms 334848 KB Output is correct
7 Correct 102 ms 336212 KB Output is correct
8 Correct 101 ms 335700 KB Output is correct
9 Correct 116 ms 335728 KB Output is correct
10 Correct 113 ms 335700 KB Output is correct
11 Correct 100 ms 335620 KB Output is correct
12 Correct 100 ms 335992 KB Output is correct

Subtask #3
30.0 / 30.0

# Verdict Execution time Memory Grader output
1 Correct 424 ms 466088 KB Output is correct
2 Correct 411 ms 465972 KB Output is correct
3 Correct 243 ms 421320 KB Output is correct
4 Correct 204 ms 420964 KB Output is correct
5 Correct 407 ms 421192 KB Output is correct
6 Correct 411 ms 421200 KB Output is correct
7 Correct 397 ms 421196 KB Output is correct

Subtask #4
30.0 / 30.0

# Verdict Execution time Memory Grader output
1 Correct 133 ms 334928 KB Output is correct
2 Correct 106 ms 334928 KB Output is correct
3 Correct 103 ms 334932 KB Output is correct
4 Correct 102 ms 334928 KB Output is correct
5 Correct 110 ms 334760 KB Output is correct
6 Correct 118 ms 334848 KB Output is correct
7 Correct 102 ms 336212 KB Output is correct
8 Correct 101 ms 335700 KB Output is correct
9 Correct 116 ms 335728 KB Output is correct
10 Correct 113 ms 335700 KB Output is correct
11 Correct 100 ms 335620 KB Output is correct
12 Correct 100 ms 335992 KB Output is correct
13 Correct 424 ms 466088 KB Output is correct
14 Correct 411 ms 465972 KB Output is correct
15 Correct 243 ms 421320 KB Output is correct
16 Correct 204 ms 420964 KB Output is correct
17 Correct 407 ms 421192 KB Output is correct
18 Correct 411 ms 421200 KB Output is correct
19 Correct 397 ms 421196 KB Output is correct
20 Correct 396 ms 421168 KB Output is correct
21 Correct 204 ms 420952 KB Output is correct
22 Correct 404 ms 421100 KB Output is correct
23 Correct 245 ms 421112 KB Output is correct
24 Correct 409 ms 421172 KB Output is correct
25 Correct 247 ms 420780 KB Output is correct
26 Correct 375 ms 466452 KB Output is correct
27 Correct 388 ms 466492 KB Output is correct