제출 #293895

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
293895		shivensinha4	Chase (CEOI17_chase)	C++17	100 / 100	1207 ms	461324 KiB

chase

#pragma GCC optimize("Ofast")
#pragma GCC target("avx,avx2,fma")
#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 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<pair<ll, int>> mxTake(v+1), mmxTake(v+1), mxNoTake(v+1), mmxNoTake(v+1);
	
	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];
			}
			if (take >= mxTake[ct].first) {
				swap(mxTake[ct], mmxTake[ct]);
				mxTake[ct] = {take, i};
			} else if (take > mmxTake[ct].first) {
				mmxTake[ct] = {take, i};
			}
			
			if (noTake >= mxNoTake[ct].first) {
				swap(mxNoTake[ct], mmxNoTake[ct]);
				mxNoTake[ct] = {noTake, i};
			} else if (noTake > mmxNoTake[ct].first) {
				mmxNoTake[ct] = {noTake, i};
			}
		}
	}
	
	ans = max({ans, mxTake[v-1].first + s, mxNoTake[v].first});
	for_(i, 0, ch) if (adj[p][i] != parent) {
		vector<ll> temp(v+1);
		for_(ct, 0, v+1) {
			ll nt = mxNoTake[ct].first, t;
			if (ct) t = mxTake[ct-1].first;
			if (mxNoTake[ct].second == adj[p][i]) nt = mmxNoTake[ct].first;
			if (ct and mxTake[ct-1].second == adj[p][i]) t = mmxTake[ct-1].first;
			
			temp[ct] = nt;
			if (ct) {
				ll nv = t + 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: 30 / 30
• Subtask 4: 30 / 30