oj.uz

Submission #922610

# Submission time Handle Problem Language Result Execution time Memory
922610 2024-02-05T19:43:10 Z NK_ Chase (CEOI17_chase) C++17
0 / 100
 278 ms 524288 KB 
// 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.0

# Verdict Execution time Memory Grader output
1 Correct 0 ms 344 KB Output is correct
2 Correct 0 ms 348 KB Output is correct
3 Correct 0 ms 348 KB Output is correct
4 Correct 0 ms 500 KB Output is correct
5 Runtime error 1 ms 604 KB Execution killed with signal 11
6 Halted 0 ms 0 KB -

Subtask #2
0 / 20.0

# Verdict Execution time Memory Grader output
1 Correct 0 ms 344 KB Output is correct
2 Correct 0 ms 348 KB Output is correct
3 Correct 0 ms 348 KB Output is correct
4 Correct 0 ms 500 KB Output is correct
5 Runtime error 1 ms 604 KB Execution killed with signal 11
6 Halted 0 ms 0 KB -

Subtask #3
0 / 30.0

# Verdict Execution time Memory Grader output
1 Runtime error 278 ms 524288 KB Execution killed with signal 9
2 Halted 0 ms 0 KB -

Subtask #4
0 / 30.0

# Verdict Execution time Memory Grader output
1 Correct 0 ms 344 KB Output is correct
2 Correct 0 ms 348 KB Output is correct
3 Correct 0 ms 348 KB Output is correct
4 Correct 0 ms 500 KB Output is correct
5 Runtime error 1 ms 604 KB Execution killed with signal 11
6 Halted 0 ms 0 KB -