Submission #293154

#	Time	Username	Problem	Language	Result	Execution time	Memory
293154		shivensinha4	Chase (CEOI17_chase)	C++17	40 / 100	1088 ms	524292 KiB

chase

#include <bits/stdc++.h> 
using namespace std; 
#define for_(i, s, e) for (int i = s; i < (int) e; i++)
#define for__(i, s, e) for (ll i = s; i < e; i++)
typedef long long ll;
typedef vector<int> vi;
typedef pair<int, int> ii;
#define endl '\n'

const int MXN = 1e5, MXV = 100;
int n, v;
ll dp[MXN+1][MXV+1];
ll nums[MXN+1];
vi adj[MXN+1];
ll ans = 0;

void prefMax(vector<ll> &k) {
	for_(i, 1, k.size()) k[i] = max(k[i], k[i-1]);
}

void sufMax(vector<ll> &k) {
	for (int i = k.size()-2; i >= 0; i--) k[i] = max(k[i], k[i+1]);
}

void init(int p, int parent) {
	ll s = 0;
	for (int i: adj[p]) if (i != parent) {
		init(i, p);
		s += nums[i];
	}
	for_(ct, 0, v+1) {
		ll take = 0, noTake = 0;
		for (int i: adj[p]) if (i != parent) {
			if (ct) take = max(take, dp[i][ct-1]);
			noTake = max(noTake, dp[i][ct]);
		}
		dp[p][ct] = noTake;
		if (ct) dp[p][ct] = max(dp[p][ct], s+take);
	}
}

void reroot(int p, int parent, vector<ll> &pVal) {
	int ch = adj[p].size();
	
	ll s = 0;
	for_(i, 0, ch) s += nums[adj[p][i]];
	vector<vector<ll>> takeVal(v+1), takePref, takeSuf, noTakeVal(v+1), noTakePref, noTakeSuf;
	
	for_(ct, 0, v+1) {
		for (int i: adj[p]) {
			ll take = 0, noTake = 0;
			if (i != parent) {
				take = dp[i][ct];
				noTake = dp[i][ct];
			} else {
				take = pVal[ct];
				noTake = pVal[ct];
			}
			takeVal[ct].push_back(take);
			noTakeVal[ct].push_back(noTake);
		}
	}
	takePref = takeSuf = takeVal; noTakePref = noTakeSuf = noTakeVal;
	takeVal.clear(); noTakeVal.clear();
	
	for_(ct, 0, v+1) {
		prefMax(takePref[ct]); prefMax(noTakePref[ct]);
		sufMax(takeSuf[ct]); sufMax(noTakeSuf[ct]);
	}
	
	ans = max({ans, takeSuf[v-1][0] + s, noTakeSuf[v][0]});
	for_(i, 0, ch) if (adj[p][i] != parent) {
		vector<ll> temp(v+1);
		for_(ct, 0, v+1) {
			temp[ct] = max(i > 0 ? noTakePref[ct][i-1] : 0, i < ch-1 ? noTakeSuf[ct][i+1] : 0);
			if (ct) {
				ll nv = max(i > 0 ? takePref[ct-1][i-1] : 0, i < ch-1 ? takeSuf[ct-1][i+1] : 0) + s - nums[adj[p][i]];
				temp[ct] = max(temp[ct], nv);
			}
			
		}
		reroot(adj[p][i], p, temp);
	}
}

int main() {
	#ifdef shiven
	freopen("test.in", "r", stdin);
	#endif
	
	ios_base::sync_with_stdio(false);
	cin.tie(0);
	
	cin >> n >> v;
	for_(i, 0, n) cin >> nums[i];
	for_(i, 0, n-1) {
		int a, b; cin >> a >> b;
		a -= 1; b -= 1;
		adj[a].push_back(b); adj[b].push_back(a);
	}
	
	if (v == 0) {
		cout << 0 << endl;
		return 0;
	}
	
	init(0, 0);
	vector<ll> useless;
	reroot(0, 0, useless);
	
	cout << ans << endl;

	return 0;
}

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