Submission #978735

#	Time	Username	Problem	Language	Result	Execution time	Memory
978735		duckindog	Chase (CEOI17_chase)	C++17	100 / 100	299 ms	170308 KiB

chase

#include <bits/stdc++.h>

using namespace std;

#define int long long

const int N = 100'000 + 10,
          B = 100 + 10;
int n, b;
int a[N];
vector<int> ad[N];

struct sortedArry { 
  pair<long long, int> arr[3];
  void add(pair<long long, int> value) { 
    arr[2] = value;
    sort(arr, arr + 3, greater<>());
  }
};

long long answer;
long long dp1[N][B], dp2[N][B];
void dfs(int u, int p = 0) { 
  long long total = 0;
  for (const auto& v : ad[u]) {
    total += a[v];
    if (v == p) continue;
    dfs(v, u);
  }

  dp2[u][1] = total;
  vector<sortedArry> best(b + 1);
  best[1].add({total, 0});
  for (const auto& v : ad[u]) { 
    if (v == p) continue;
    for (int i = 1; i <= b; ++i) { 
      auto&ret = dp1[u][i];
      ret = max({ret, dp1[v][i], dp1[v][i - 1] + total - a[p]});
    }
    for (int i = 1; i <= b; ++i) { 
      auto& ret = dp2[u][i];
      long long val = max(dp2[v][i], dp2[v][i - 1] + total - a[v]);
      ret = max({ret, val});
      best[i].add({val, v});
    }
  }
  
  for (const auto& v : ad[u]) { 
    if (v == p) continue;
    for (int i = 1; i <= b; ++i) { 
      for (int j = 0; j < 2; ++j) { 
        const auto& [value, x] = best[i].arr[j];
        if (x == v) continue;
        answer = max(answer, value + dp1[v][b - i]);
      }
    }
  }
}

int32_t main() { 
  cin.tie(0)->sync_with_stdio(0);

  cin >> n >> b;
  for (int i = 1; i <= n; ++i) cin >> a[i];
  for (int i = 2; i <= n; ++i) {
    int u, v; cin >> u >> v;
    ad[u].push_back(v);
    ad[v].push_back(u);
  }
  
  dfs(1);

  cout << answer << "\n";
}

• Subtask 1: 20 / 20
• Subtask 2: 20 / 20
• Subtask 3: 30 / 30
• Subtask 4: 30 / 30