oj.uz

Submission #941785

# Submission time Handle Problem Language Result Execution time Memory
941785 2024-03-09T12:07:47 Z vjudge1 Chase (CEOI17_chase) C++17
100 / 100
 393 ms 466516 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 76 ms 334744 KB Output is correct
2 Correct 80 ms 334676 KB Output is correct
3 Correct 79 ms 334952 KB Output is correct
4 Correct 77 ms 334852 KB Output is correct
5 Correct 77 ms 334888 KB Output is correct
6 Correct 77 ms 334888 KB Output is correct

Subtask #2
20.0 / 20.0

# Verdict Execution time Memory Grader output
1 Correct 76 ms 334744 KB Output is correct
2 Correct 80 ms 334676 KB Output is correct
3 Correct 79 ms 334952 KB Output is correct
4 Correct 77 ms 334852 KB Output is correct
5 Correct 77 ms 334888 KB Output is correct
6 Correct 77 ms 334888 KB Output is correct
7 Correct 79 ms 336212 KB Output is correct
8 Correct 78 ms 335692 KB Output is correct
9 Correct 80 ms 335700 KB Output is correct
10 Correct 82 ms 335696 KB Output is correct
11 Correct 80 ms 335700 KB Output is correct
12 Correct 78 ms 335744 KB Output is correct

Subtask #3
30.0 / 30.0

# Verdict Execution time Memory Grader output
1 Correct 393 ms 463956 KB Output is correct
2 Correct 370 ms 463628 KB Output is correct
3 Correct 203 ms 419012 KB Output is correct
4 Correct 229 ms 419120 KB Output is correct
5 Correct 352 ms 419156 KB Output is correct
6 Correct 340 ms 418900 KB Output is correct
7 Correct 340 ms 419112 KB Output is correct

Subtask #4
30.0 / 30.0

# Verdict Execution time Memory Grader output
1 Correct 76 ms 334744 KB Output is correct
2 Correct 80 ms 334676 KB Output is correct
3 Correct 79 ms 334952 KB Output is correct
4 Correct 77 ms 334852 KB Output is correct
5 Correct 77 ms 334888 KB Output is correct
6 Correct 77 ms 334888 KB Output is correct
7 Correct 79 ms 336212 KB Output is correct
8 Correct 78 ms 335692 KB Output is correct
9 Correct 80 ms 335700 KB Output is correct
10 Correct 82 ms 335696 KB Output is correct
11 Correct 80 ms 335700 KB Output is correct
12 Correct 78 ms 335744 KB Output is correct
13 Correct 393 ms 463956 KB Output is correct
14 Correct 370 ms 463628 KB Output is correct
15 Correct 203 ms 419012 KB Output is correct
16 Correct 229 ms 419120 KB Output is correct
17 Correct 352 ms 419156 KB Output is correct
18 Correct 340 ms 418900 KB Output is correct
19 Correct 340 ms 419112 KB Output is correct
20 Correct 344 ms 420296 KB Output is correct
21 Correct 146 ms 421148 KB Output is correct
22 Correct 390 ms 421176 KB Output is correct
23 Correct 161 ms 420944 KB Output is correct
24 Correct 363 ms 421204 KB Output is correct
25 Correct 191 ms 420952 KB Output is correct
26 Correct 322 ms 466516 KB Output is correct
27 Correct 324 ms 466480 KB Output is correct