oj.uz

Submission #922653

# Submission time Handle Problem Language Result Execution time Memory
922653 2024-02-05T21:48:51 Z NK_ Chase (CEOI17_chase) C++17
0 / 100
 410 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];
	}

	ll ans = 0;
	V<V<vl>> up(N, V<vl>(K + 1, vl(2, 0))), dwn(N, V<vl>(K + 1, vl(2, 0)));
	function<void(int, int)> gen = [&](int u, int p) {

		if (K) up[u][1][1] = dwn[u][1][1] = P[u];

		vl mx0(K + 1), mx1(K + 1);
		for(auto& v : adj[u]) if (v != p) {
			gen(v, u);	

			for(int k = 0; k <= K; k++) {
				// cout << u << " " << v << " " << k << " => 0 => " << mx0[k] << " " << up[u][K - k][0] << endl;
				// if (k) cout << u << " " << v << " " << k << " => 1 => " << mx1[k] << " " << up[u][K + 1 - k][1] - P[u] << endl;

				ans = max(ans, mx0[k] + up[u][K - k][0]);
				if (k) ans = max(ans, mx1[k] + up[u][K + 1 - k][1] - P[u]);
			}

			for(int k = 0; k <= K; k++) {
				mx0[k] = max(mx0[k], dwn[u][k][0]);
				mx1[k] = max(mx1[k], dwn[u][k][1]);

				if (k) mx0[k] = max(mx0[k], mx0[k - 1]), mx0[k] = max(mx0[k], mx0[k - 1]);

				up[u][k][0] = max({up[u][k][0], up[v][k][0], up[v][k][1]});
				if (k) up[u][k][1] = max({up[u][k][1], up[v][k - 1][0] + P[u] - A[v], up[v][k - 1][1] + P[u] - A[v]});

				dwn[u][k][0] = max({dwn[u][k][0], dwn[v][k][1] - A[u], dwn[v][k][0]});
				if (k) dwn[u][k][1] = max({dwn[u][k][1], dwn[v][k - 1][1] - A[u] + P[u], dwn[v][k - 1][0] + P[u]});
			}
		}

		for(int k = 0; k <= K; k++) ans = max({ans, up[u][k][0], up[u][k][1], dwn[u][k][0], dwn[u][k][1]});

		// cout << "NODE: " << u << " |  CUR ANS: " << ans << endl;
		// for(int k = 0; k <= K; k++) for(int x = 0; x < 2; x++) cout << k << " " << x << " => " << up[u][k][x] << " " << dwn[u][k][x] << endl;

		// for(int x = 0; x < 2; x++) {
		// 	vl mx(K + 1); mx[0] = up[u][0][x]; 
		// 	for(int i = 1; i <= K; i++) mx[i] = max(up[u][i][x], mx[i - 1]);
		// 	for(int i = 0; i <= K; i++) {
		// 		cout << i << " => " << mx[K - i] << " " << dwn[u][i][x] << " " << x * P[u] << endl;
		// 		ans = max(ans, mx[K - i] + dwn[u][i][x] - x * P[u]);
		// 	}
		// }
	};
	
	gen(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 Incorrect 0 ms 500 KB Output isn't correct
3 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 Incorrect 0 ms 500 KB Output isn't correct
3 Halted 0 ms 0 KB -

Subtask #3
0 / 30.0

# Verdict Execution time Memory Grader output
1 Runtime error 410 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 Incorrect 0 ms 500 KB Output isn't correct
3 Halted 0 ms 0 KB -