Submission #649048

#	Time	Username	Problem	Language	Result	Execution time	Memory
649048		dutinmeow	Transport (COCI19_transport)	C++17	130 / 130	640 ms	28248 KiB

Main

#include <bits/stdc++.h>

#include <ext/pb_ds/assoc_container.hpp> 
#include <ext/pb_ds/tree_policy.hpp>
 
template<class K, class V>
using ordered_map = __gnu_pbds::tree<K, V, std::less<K>, __gnu_pbds::rb_tree_tag, __gnu_pbds::tree_order_statistics_node_update>;
 
template<class K>
using ordered_set = ordered_map<K, __gnu_pbds::null_type>;

namespace std {

template<class Fun>
class y_combinator_result {
	Fun fun_;
public:
	template<class T>
	explicit y_combinator_result(T &&fun): fun_(std::forward<T>(fun)) {}

	template<class ...Args>
	decltype(auto) operator()(Args &&...args) {
		return fun_(std::ref(*this), std::forward<Args>(args)...);
	}
};

template<class Fun>
decltype(auto) y_combinator(Fun &&fun) {
	return y_combinator_result<std::decay_t<Fun>>(std::forward<Fun>(fun));
}

} // namespace std

int main() {
	using namespace std;
	ios_base::sync_with_stdio(false);
	cin.tie(NULL);

	int N;
	cin >> N;
	vector<long long> A(N);
	for (auto &a : A)
		cin >> a;
	vector<vector<pair<int, long long>>> T(N);
	for (int i = 1; i < N; i++) {
		int u, v;
		long long w;
		cin >> u >> v >> w;
		u--, v--;
		T[u].emplace_back(v, w);
		T[v].emplace_back(u, w);
	}

	vector<bool> blk(N, false);
	vector<int> wei(N);

	auto init_wei = y_combinator([&](auto self, int u, int p = -1) -> int {
		wei[u] = 1;
		for (auto [v, w] : T[u]) {
			if (v == p || blk[v])
				continue;
			wei[u] += self(v, u);
		}
		return wei[u];
	});

	auto find_cen = y_combinator([&](auto self, int n, int u, int p = -1) -> int {
		for (auto [v, w] : T[u]) {
			if (v == p || blk[v])
				continue;
			if (2 * wei[v] > n)
				return self(n, v, u);
		}
		return u;
	});

	// compute fuel from u to root
	auto dfs0 = y_combinator([&](auto self, long long tot, long long mgp, vector<long long> &src, int u, int p) -> void {
		if (mgp >= 0) 
			src.push_back(tot);
		for (auto [v, w] : T[u]) {
			if (v == p || blk[v])
				continue;
			auto ntot = tot - w + A[v];
			auto nmgp = min(mgp, 0LL) - w + A[v];
			self(ntot, nmgp, src, v, u);
		}
	});

	// compute fuel from root to u
	auto dfs1 = y_combinator([&](auto self, long long cur, long long ret, vector<long long> &snk, int u, int p) -> void {
		snk.push_back(ret);
		cur += A[u];
		for (auto [v, w] : T[u]) {
			if (v == p || blk[v])
				continue;
			if (cur < w)
				self(0, ret + w - cur, snk, v, u);
			else
				self(cur - w, ret, snk, v, u);
		}
	});

	cout << y_combinator([&](auto self, int u = 0) -> long long {
		int n = init_wei(u);
		u = find_cen(n, u);

		int ind = 0;
		long long ret = 0;
		ordered_set<pair<long long, int>> ord1, ord2;
		ord1.insert(make_pair(0, ind++));
		ord2.insert(make_pair(0, ind++));
		for (auto [v, w] : T[u]) {
			if (blk[v])
				continue;
			vector<long long> src, snk;
			dfs0(A[v] - w, A[v] - w, src, v, u);
			dfs1(max(A[u] - w, 0LL), max(w - A[u], 0LL), snk, v, u);
			long long tem = 0;
			for (auto x : src) 
				tem += ord2.order_of_key(make_pair(x, INT_MAX));
			for (auto x : snk)
				tem += (int)ord1.size() - ord1.order_of_key(make_pair(x, INT_MIN));
			for (auto x : src)
				ord1.insert(make_pair(x, ind++));
			for (auto x : snk)
				ord2.insert(make_pair(x, ind++));
			ret += tem;
		}

		blk[u] = true;
		for (auto [v, w] : T[u])
			if (!blk[v])
				ret += self(v);
		return ret;
	})() << '\n';
}

• Subtask 1: 13 / 13
• Subtask 2: 13 / 13
• Subtask 3: 13 / 13
• Subtask 4: 13 / 13
• Subtask 5: 13 / 13
• Subtask 6: 13 / 13
• Subtask 7: 13 / 13
• Subtask 8: 13 / 13
• Subtask 9: 13 / 13
• Subtask 10: 13 / 13