Submission #888245

#TimeUsernameProblemLanguageResultExecution timeMemory
888245MarkadiuszChase (CEOI17_chase)C++17
100 / 100
514 ms343380 KiB
#include<bits/stdc++.h>
using namespace std;
using LL=long long;
#define FOR(i,l,r) for(int i=(l);i<=(r);++i)
#define REP(i,n) FOR(i,0,(n)-1)
#define ssize(x) int(x.size())
template<class A,class B>auto&operator<<(ostream&o,pair<A,B>p){return o<<'('<<p.first<<", "<<p.second<<')';}
template<class T>auto operator<<(ostream&o,T x)->decltype(x.end(),o){o<<'{';int i=0;for(auto e:x)o<<(", ")+2*!i++<<e;return o<<'}';}
#ifdef DEBUG
#define debug(x...) cerr<<"["#x"]: ",[](auto...$){((cerr<<$<<"; "),...);}(x),cerr<<'\n'
#else
#define debug(...) {}
#endif

int main() {
	cin.tie(0)->sync_with_stdio(0);

	int n, V;
	cin >> n >> V;
	if (V == 0) {
		cout << 0 << '\n';
		return 0;
	}
	vector<int> P(n);
	REP(i, n)
		cin >> P[i];
	vector<vector<int>> graph(n);
	vector<LL> sum_neighbours(n);
	REP(i, n - 1) {
		int a, b;
		cin >> a >> b;
		--a, --b;
		graph[a].emplace_back(b);
		graph[b].emplace_back(a);
		sum_neighbours[a] += P[b];
		sum_neighbours[b] += P[a];
	}
	debug(n, V, P, graph, sum_neighbours);

#ifdef LOCAL
	const LL INF = 100;
#else
	const LL INF = 1e18;
#endif
	using Info = vector<array<LL, 2>>;
	Info blank(V + 1);
	FOR(i, 0, V) {
		blank[i][0] = blank[i][1] = 0;
	}
	blank[0][1] = -INF;
	auto make_info = [&](int v) {
		Info ret = blank;
		ret[1][1] = sum_neighbours[v];
		return ret;
	};
	auto extend_info = [&](const Info& h, int v) -> Info {
		Info ret = blank;
		FOR(i, 0, V) {
			ret[i][0] = max(h[i][0], h[i][1] - P[v]);
		}
		FOR(i, 1, V) {
			ret[i][1] = ret[i - 1][0] + sum_neighbours[v];
		}
		debug(h, v, P[v], ret);
		return ret;
	};
	auto get_better = [&](Info a, const Info& b) {
		FOR(i, 0, V)
			REP(j, 2)
				a[i][j] = max(a[i][j], b[i][j]);
		return a;
	};
	auto get_ans = [&](const Info& a) {
		LL ans = 0;
		FOR(i, 0, V)
			REP(j, 2)
				ans = max(ans, a[i][j]);
		return ans;
	};

	vector<Info> best(n);
	function<void(int, int)> init_dfs = [&](int v, int p) {
		auto it = find(graph[v].begin(), graph[v].end(), p);
		if (it != graph[v].end())
			graph[v].erase(it);
		debug(v, p, graph[v]);

		for (auto x : graph[v])
			init_dfs(x, v);

		auto val = make_info(v);
		for (int x : graph[v]) {
			val = get_better(val, extend_info(best[x], v));
		}
		best[v] = val;
	};
	init_dfs(0, -1);

	LL ans = 0;
	function<void(int, const Info&)> dfs = [&](int v, const Info& value_from_parent) {
		const int s = ssize(graph[v]);
		vector<Info> values_for_children;
		{
			const auto extended_value_from_parent = extend_info(value_from_parent, v);
			ans = max(ans, get_ans(get_better(extended_value_from_parent, best[v])));
			values_for_children.resize(s, extended_value_from_parent);
		}
		{
			auto val = make_info(v);
			REP(i, s) {
				values_for_children[i] = get_better(values_for_children[i], val);
				const int x = graph[v][i];
				val = get_better(move(val), extend_info(best[x], v));
			}
		}
		{
			auto val = make_info(v);
			for (int i = s - 1; i >= 0; --i) {
				values_for_children[i] = get_better(values_for_children[i], val);
				const int x = graph[v][i];
				val = get_better(move(val), extend_info(best[x], v));
			}
		}
		REP(i, ssize(graph[v]))
			dfs(graph[v][i], move(values_for_children[i]));
	};
	dfs(0, blank);
	cout << ans << '\n';

#ifdef LOCAL
	system("grep VmPeak /proc/$PPID/status >&2");
#endif
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...