제출 #922610

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
922610		NK_	Chase (CEOI17_chase)	C++17	0 / 100	278 ms	524288 KiB

chase

// Success consists of going from failure to failure without loss of enthusiasm
#include <bits/stdc++.h>

using namespace std;

#define nl '\n'
#define pb push_back
#define sz(x) int(x.size())

using ll = long long;
template<class T> using V = vector<T>;
using vl = V<ll>;
using vi = V<int>;

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

	int N, K; cin >> N >> K;
	vi A(N); for(auto& x : A) cin >> x;

	vl P(N);
	V<vi> adj(N); for(int i = 0; i < N - 1; i++) {
		int u, v; cin >> u >> v; --u, --v;
		adj[u].pb(v);
		adj[v].pb(u);
		P[u] += A[v], P[v] += A[u];
	}

	V<vl> dp(N, vl(K+1, 0)); 
	V<V<multiset<ll>>> S(N, V<multiset<ll>>(K+1));

	for(int i = 0; i < N; i++) S[i][0].insert(0), S[i][1].insert(P[i]);


	auto add = [&](int u, int v) { // add v as a child of u
		for(int i = 0; i <= K; i++) {
			S[u][i].insert(dp[v][i]);
			if (i + 1 <= K) S[u][i + 1].insert(dp[v][i] - A[v] + P[u]);
		}
	};

	auto rem = [&](int u, int v) { // remove v as a child of u
		for(int i = 0; i <= K; i++) {
			S[u][i].erase(S[u][i].find(dp[v][i]));
			if (i + 1 <= K) S[u][i + 1].erase(S[u][i + 1].find(dp[v][i] - A[v] + P[u]));
		}
	};

	auto calc = [&](int u) {
		for(int i = 0; i <= K; i++) dp[u][i] = (sz(S[u][i]) ? *rbegin(S[u][i]) : 0);
	};

	function<void(int, int)> gen = [&](int u, int p) {
		for(auto& v : adj[u]) if (v != p) {
			gen(v, u);
			add(u, v);
		}

		calc(u);
	};
	
	gen(0, -1);


	ll ans = 0;
	function<void(int, int)> dfs = [&](int u, int p) {
		ans = max(ans, *max_element(begin(dp[u]), end(dp[u])));

		for(auto& v : adj[u]) if (v != p) {
			rem(u, v); calc(u);
			add(v, u); calc(v);

			dfs(v, u);

			rem(v, u); calc(v);
			add(u, v); calc(u);
		}
	};

	dfs(0, -1);

	cout << ans << nl;

	exit(0-0);
}

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