Submission #97746

#	Time	Username	Problem	Language	Result	Execution time	Memory
97746		Anai	Construction of Highway (JOI18_construction)	C++14	100 / 100	939 ms	28936 KiB

construction

#include <bits/stdc++.h>
#define x first
#define y second
using namespace std;

using i64 = long long;
using pii = pair<int, int>;

const int N = 1e5 + 5;


set<pii> paths[N];
vector<int> g[N];
i64 aib[N];
int cost[N], chain[N], top[N], far[N], siz[N], h[N];

vector<pii> edges;
int n, chain_ptr;


static void dfs(int u) {
	siz[u] = 1;
	if (g[u].empty()) {
		chain[u] = ++chain_ptr;
		top[chain[u]] = u;
		return; }

	for (auto v: g[u]) {
		far[v] = u;
		h[v] = h[u] + 1;
		dfs(v);
		siz[u]+= siz[v]; }

	for (auto &v: g[u])
		if (siz[v] > siz[g[u][0]])
			swap(v, g[u][0]);
	chain[u] = chain[g[u][0]];
	top[chain[u]] = u; }

static void normalize() {
	map<int, int> mp;
	int ptr(0);

	for (int i = 1; i <= n; ++i)
		mp[cost[i]] = 0;
	for (auto &i: mp)
		i.second = ++ptr;
	for (int i = 1; i <= n; ++i)
		cost[i] = mp[cost[i]]; }

static int lsb(const int &arg) {
	return arg & -arg; }

static void update(int pos, int val) {
	for (; pos < N; pos+= lsb(pos))
		aib[pos]+= val; }

static i64 query(int pos) {
	i64 ant = 0;
	for (; pos > 0; pos-= lsb(pos))
		ant+= aib[pos];
	return ant; }

static i64 getcost(int nod) {
	vector<pii> v;
	i64 total, ant;
	int c;

	c = chain[nod = far[nod]];
	while (c) {
		auto it = paths[c].lower_bound(pii(h[nod], -1));
		if (it != end(paths[c]))
			v.emplace_back(cost[it->second], h[nod]);

		while (it != begin(paths[c])) {
			it = prev(it);
			v.emplace_back(cost[it->second], it->first); }

		nod = far[top[c]];
		c = chain[nod]; }

	for (int i = 0; i < v.size() - 1; ++i)
		v[i].second-= v[i + 1].second;
	v.back().second+= 1;


	ant = 0;
	for (auto i: v) {
		ant+= query(i.first - 1) * i.second;
		update(i.first, i.second); }
	for (auto i: v) {
		update(i.first, -i.second); }

	return ant; }

static void enforce(int nod) {
	int origin = nod;
	int c = chain[nod];

	while (c) {
		set<pii>::iterator tmp, it = paths[c].upper_bound(pii(h[nod], n));

		if (it != begin(paths[c])) {
			it = prev(it);
			while (true) {
				if (it == begin(paths[c])) {
					paths[c].erase(it);
					break; }
				else {
					tmp = prev(it);
					paths[c].erase(it);
					it = tmp; } } }

		paths[c].emplace(h[nod], origin);
		nod = far[top[c]];
		c = chain[nod]; } }

int main() {
#ifdef HOME
	freopen("joi_construction.in", "r", stdin);
	freopen("joi_construction.out", "w", stdout);
#endif
	ios::sync_with_stdio(false);
	cin.tie(0), cout.tie(0);

	cin >> n;
	for (int i = 1; i <= n; ++i)
		cin >> cost[i];

	for (int a, b, i = 1; i < n; ++i) {
		cin >> a >> b;
		edges.emplace_back(a, b);
		g[a].push_back(b); }

	normalize();
	dfs(1);
	paths[chain[1]].emplace(h[1], 1);

	for (const auto &e: edges) {
		cout << getcost(e.second) << '\n';
		enforce(e.second); }

	return 0; }

Compilation message (stderr)

construction.cpp: In function 'i64 getcost(int)':
construction.cpp:82:20: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
  for (int i = 0; i < v.size() - 1; ++i)
                  ~~^~~~~~~~~~~~~~
construction.cpp:66:6: warning: unused variable 'total' [-Wunused-variable]
  i64 total, ant;
      ^~~~~

• Subtask 1: 7 / 7
• Subtask 2: 9 / 9
• Subtask 3: 84 / 84