Submission #1000469

#	Time	Username	Problem	Language	Result	Execution time	Memory
1000469		0npata	Chase (CEOI17_chase)	C++17	0 / 100	190 ms	524288 KiB

chase

#include<bits/stdc++.h>
using namespace std;
 
#define int long long
#define vec vector
#define fst first
#define snd second
#define arr array
 
const int N = 100'005;
const int K = 105;
 
int vals[N];
arr<pair<int, int>, 2> bsts[N][K][2];
vec<int> tree[N];
 
void merge(int u, int k, bool sel, pair<int, int> res) {
	if(res.fst >= bsts[u][k][sel][0].fst) {
		if(res.snd != bsts[u][k][sel][0].snd) {
			bsts[u][k][sel][1] = bsts[u][k][sel][0];
		}
		bsts[u][k][sel][0] = res;
	}
	else if(res.fst >= bsts[u][k][sel][1].fst && res.snd != bsts[u][k][sel][0].snd) {
		bsts[u][k][sel][1] = res;
	}
}
 
// given two vertices u and v, compute for u, the result if the path started at u and went into v
// compute the result for all possible cases of k remaining selections or whether the node is selected
void upd(int u, int v, int usum) {
	// for each k, we will try to merge the new result with bsts[u][k][..]
	// if there are no selections remaining, then there is no point to going to edge v
	merge(u, 0, false, {0, v});
	merge(u, 0, true, {usum-vals[u], v});
	for(int k = 1; k<K; k++) {
		int vres = bsts[v][k][false][0].snd == u ? bsts[v][k][false][1].fst : bsts[v][k][false][0].fst;
 
		int ures_nsel = vres;
		merge(u, k, false, {ures_nsel, v});
		int ures_sel = vres + usum - vals[u];
		merge(u, k, true, {ures_sel, v});
	}
	for(int k = 1; k<K; k++) {
		int vres = bsts[v][k-1][true][0].snd == u ? bsts[v][k-1][true][1].fst : bsts[v][k-1][true][0].fst;
 
		int ures_nsel = vres - vals[u];
		merge(u, k, false, {ures_nsel, v});
		int ures_sel = vres + usum - vals[u] - vals[u];
		merge(u, k, true, {ures_sel, v});
	}
}
 
void dfs1(int u, int p=-1) {
	int usum = vals[u];
	for(int v : tree[u]) {
		usum += vals[v];
 
		if(v == p) continue;
 
		dfs1(v, u);
	}
	for(int i = 0; i<K; i++){
		merge(u, i, false, {0, -1});
		merge(u, i, true, {usum-vals[u], -1});
	}
 
 
	for(int v : tree[u]) {
		upd(u, v, usum);
	}
}
 
void dfs2(int u, int p = -1) {
	int usum = vals[u];
	for(int v : tree[u]) usum += vals[v];
 
	if(p != -1) {
		upd(u, p, usum);
	}
 
	for(int v : tree[u]) {
		if(v==p) continue;
		dfs2(v, u);
	}
}
 
int32_t main() {
	for(int i = 0; i<N; i++) {
		for(int j = 0; j<K; j++) {
			for(int k = 0; k<2; k++) {
				for(int l = 0; l<2; l++) {
					bsts[i][j][k][l] = {-1e9, -1};
				}
			}
		}
	}
 
	ios_base::sync_with_stdio(false);
	cin.tie(0);
 
	int n, k;
	cin >> n >> k;
 
	for(int i = 0; i<n; i++) cin >> vals[i];
 
	for(int i = 0; i<n-1; i++) {
		int u, v;
		cin >> u >> v;
		u--; v--;
 
		tree[u].push_back(v);
		tree[v].push_back(u);
	}
 
	dfs1(0);
 
 
	dfs2(0);
 
	//for(int i = 0; i<n; i++) {
		//cerr << "I: " << i << " " << bsts[i][k][false][0].fst << ' ' << bsts[i][k-1][true][0].fst << '\n';
	//}
 
 
	int ans = 0;
 
	for(int i = 0; i<n; i++) {
 
		ans = max(ans, bsts[i][k][false][0].fst);
		if(k>0) ans = max(ans, bsts[i][k-1][true][0].fst);
	}
 
	cout << ans << '\n';
}

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