Submission #1261623

#	Time	Username	Problem	Language	Result	Execution time	Memory
1261623		kaiboy	Chase (CEOI17_chase)	C++20	100 / 100	309 ms	248028 KiB

chase

#include <algorithm>
#include <iostream>
#include <vector>

using namespace std;

const       int   N = 100000;
const       int   K = 100;
const long long INF = 0x3f3f3f3f3f3f3f3fLL;

int aa[N], *ej[N], eo[N], k_;
long long bb[N], dp[N][K + 1][2], dq[N][K + 1], xx[K + 1];

void append(int i, int j) {
	int o = eo[i]++;
	if (!o)
		ej[i] = (int *) malloc(sizeof *ej[i]);
	else if (!(o & o - 1))
		ej[i] = (int *) realloc(ej[i], (o << 1) * sizeof *ej[i]);
	ej[i][o] = j;
}

void detach(int i, int j) {
	for (int o = 0; o < eo[i]; o++)
		if (ej[i][o] == j) {
			swap(ej[i][o], ej[i][--eo[i]]);
			return;
		}
}

void dfs0(int p, int i) {
	detach(i, p);
	for (int k = 0; k <= k_; k++)
		dp[i][k][0] = 0, dp[i][k][1] = k ? bb[i] : -INF;
	for (int o = 0; o < eo[i]; o++) {
		int j = ej[i][o];
		dfs0(i, j);
		for (int k = 0; k <= k_; k++) {
			dp[i][k][0] = max(dp[i][k][0], max(dp[j][k][0], dp[j][k][1]));
			if (k < k_)
				dp[i][k + 1][1] = max(dp[i][k + 1][1], max(dp[j][k][0], dp[j][k][1]) + bb[i] - aa[j]);
		}
	}
}

void dfs1(int p, int i) {
	if (p != -1)
		for (int k = 0; k <= k_; k++)
			dq[i][k] = max(0LL, k ? bb[i] - aa[p] : -INF);
	for (int o = 0; o < eo[i]; o++) {
		int j = ej[i][o];
		dfs1(i, j);
		if (p != -1)
			for (int k = 0; k <= k_; k++)
				dq[i][k] = max(dq[i][k], max(dq[j][k], k ? dq[j][k - 1] + bb[i] - aa[p] : -INF));
	}
}

int main() {
	ios_base::sync_with_stdio(false), cin.tie(NULL);
	int n; cin >> n >> k_;
	for (int i = 0; i < n; i++)
		cin >> aa[i];
	for (int h = 0; h < n - 1; h++) {
		int i, j; cin >> i >> j, i--, j--;
		append(i, j), bb[i] += aa[j];
		append(j, i), bb[j] += aa[i];
	}
	dfs0(-1, 0);
	dfs1(-1, 0);
	long long ans = 0;
	for (int i = 0; i < n; i++) {
		for (int f = 0; f < 2; f++)
			ans = max(ans, dp[i][k_][f]);
		for (int o = 0; o < eo[i]; o++) {
			int j = ej[i][o];
			ans = max(ans, max(dq[j][k_], k_ ? dq[j][k_ - 1] + bb[i] : -INF));
		}
	}
	for (int i = 0; i < n; i++) {
		for (int k = 0; k <= k_; k++)
			xx[k] = -INF;
		for (int o = 0; o < eo[i]; o++) {
			int j = ej[i][o];
			for (int k = 0; k <= k_; k++) {
				ans = max(ans, max(dp[j][k][0], dp[j][k][1]) + xx[k_ - k]);
				if (k < k_)
					ans = max(ans, max(dp[j][k][0], dp[j][k][1]) + bb[i] - aa[j] + xx[k_ - k - 1]);
			}
			for (int k = 0; k <= k_; k++)
				xx[k] = max(xx[k], dq[j][k]);
		}
		for (int k = 0; k <= k_; k++)
			xx[k] = -INF;
		for (int o = eo[i] - 1; o >= 0; o--) {
			int j = ej[i][o];
			for (int k = 0; k <= k_; k++) {
				ans = max(ans, max(dp[j][k][0], dp[j][k][1]) + xx[k_ - k]);
				if (k < k_)
					ans = max(ans, max(dp[j][k][0], dp[j][k][1]) + bb[i] - aa[j] + xx[k_ - k - 1]);
			}
			for (int k = 0; k <= k_; k++)
				xx[k] = max(xx[k], dq[j][k]);
		}
	}
	cout << ans << '\n';
	return 0;
}

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