제출 #941785

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
941785		vjudge1	Chase (CEOI17_chase)	C++17	100 / 100	393 ms	466516 KiB

chase

#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 / 20
• Subtask 2: 20 / 20
• Subtask 3: 30 / 30
• Subtask 4: 30 / 30