Submission #1345693

#TimeUsernameProblemLanguageResultExecution timeMemory
1345693limitsSplit the Attractions (IOI19_split)C++20
0 / 100
581 ms1114112 KiB
#pragma GCC optimize("Ofast")
#pragma GCC optimize("O3,unroll-loops")

#include <bits/stdc++.h>

using namespace std;

#define f0r(i, n) for (auto i = 0; i < (n); ++i)
#define fnr(i, n, k) for (auto i = (n); i < (k); ++i)
#define all(v) (v).begin(), (v).end()
#define pb push_back
#define F first
#define S second
#define ctn(x) cout << x << '\n'
#define forl(a, l) for (auto a : l)
#define ctl(l) for (auto &a : (l)) cout << a << ' '; cout << endl;
#define lb(v, x) (lower_bound(all(v), x) - begin(v))
#define ub(v, x) (upper_bound(all(v), x) - begin(v))
#define pq priority_queue

template <class T>
using V = vector<T>;
using ll = long long;
using vi = V<int>;
using vl = V<ll>;
using pi = pair<int, int>;
using ti = tuple<int, int, int>;
using Adj = V<vi>;
using vvi = V<vi>;

#include "split.h"

using namespace std;

Adj G;
vi res, sz;
int rm;

void dfs(int v, int p = -1) {
	if (v == rm) return;
	sz[v] = 1;
	forl(c, G[v]) if (c != p && c != rm) {
		dfs(c, v);
		sz[v] += sz[c];
	}
}

void get_res(int v, vi &out, int b, int p = -1) {
	if (v == rm) return;
	if (out.size() == b) {
		return;
	}
	out.pb(v);
	forl(c, G[v]) if (c != p) get_res(c, out, b, v);
}

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
	G.assign(n, {});
	res.assign(n, 0);
	sz.assign(n, 0);

	int m = p.size();
	f0r(i, m) {
		G[p[i]].pb(q[i]);
		G[q[i]].pb(p[i]);
	}


	V<pi> lab{{a, 1}, {b, 2}, {c, 3}};
	sort(all(lab));


	bool fnd = 0;

	f0r(i, 2) {
		swap(lab[0], lab[1]);
		res.assign(n, 0);
		sz.assign(n, 0);
		rm = -1;

		dfs(0);

		//vi s = sz;
		//sort(all(s));
		//s.erase(unique(all(s)), end(s));
		//ctl(s);

		V<pi> ps;
		f0r(i, n) ps.pb({sz[i], i});
		sort(all(ps));
		auto it = lower_bound(all(ps), make_pair(lab[1].F, 0));
		//cout << lab[1].F << endl;
		if (it == ps.end()) continue;
		vi out;
		get_res(it->S, out, lab[1].F);
		forl(x, out) res[x] = lab[1].S;
		rm = it->S;

		//	cout << "REM: " << it->F << endl;
		//cout << "DONE 1 " << endl;

		sz.assign(n, 0);
		dfs(0);
		ps.clear();
		f0r(i, n) ps.pb({sz[i], i});
		sort(all(ps));
		it = lower_bound(all(ps), make_pair(lab[0].F, 0));
		//	cout << "NEED: " << lab[0].F << endl;
		//	cout << "NEED: " << *max_element(all(sz)) << endl;
		if (it == ps.end()) continue;
		out.clear();
		get_res(it->S, out, lab[0].F);
		forl(x, out) res[x] = lab[0].S;

		//cout << "DONE 2 " << endl;

		f0r(i, n) if (!res[i]) res[i] = lab[2].S;

		return res;

	}

	if (!fnd) return vi(n);


	//vi cnt(3);
	//f0r(i, n) cnt[res[i]-1]++;
	//ctl(cnt);

	return res;

}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...