Submission #836566

#	Time	Username	Problem	Language	Result	Execution time	Memory
836566		penguinman	Split the Attractions (IOI19_split)	C++17	40 / 100	71 ms	17612 KiB

split

#include "split.h"
#include <bits/stdc++.h>

#ifndef EVAL
#include "grader.cpp"
#endif

using std::cin;
using std::cout;
using std::endl;
using std::vector;
using std::string;
using ll = long long;
using vi = vector<ll>;
using vii = vector<vi>;
using pii = std::pair<ll,ll>;

#define rep(i,j,k) for(ll i=ll(j); i<ll(k); i++)
#define REP(i,j,k) for(ll i=ll(j); i<=ll(k); i++)
#define per(i,j,k) for(ll i=ll(j); i>=ll(k); i--)
#define all(a) a.begin(),a.end()
#define pb emplace_back
#define mp std::make_pair
#define mtp std::make_tuple

constexpr ll inf = (1ll<<60);

struct union_find{
	ll N;
	vi par, sz;
	union_find(int n): N(n){
		par.resize(N);
        sz.resize(N,1);
		rep(i,0,N) par[i] = i;
	}
	ll root(ll now){
		if(par[now] == now) return now;
		return par[now] = root(par[now]);
	}
	bool same(ll l, ll r){
		return root(l) == root(r);
	}
	void merge(ll l, ll r){
		l = root(l);
		r = root(r);
		if(l == r) return;
		par[l] = r;
        sz[r] += sz[l];
        sz[l] = 0;
	}
    ll size(ll now){
        now = root(now);
        return sz[now];
    }
};

std::vector<int> find_split(int n, int a, int b, int c, std::vector<int> p, std::vector<int> q) {
	ll N = n;

	vector<pii> data = {mp(a,1), mp(b,2), mp(c,3)};
	sort(all(data));

	vii edge(N);
	union_find tree(N);
    vi u, v;
	rep(i,0,p.size()){
		if(tree.same(p[i], q[i])){
            u.pb(p[i]);
            v.pb(q[i]);
            continue;
        }
		tree.merge(p[i], q[i]);
		edge[p[i]].pb(q[i]);
		edge[q[i]].pb(p[i]);
	}

	vector<int> ans(N, data[2].second);

	auto modify_all = [&](){
		rep(i,0,2){
			rep(j,0,N){
				if(ans[j] == data[i].second){
					vector<bool> visited(N);
					ll cnt = 0;
					std::function<void(ll)> modify = [&](ll now){
						if(cnt < data[i].first){
							cnt++;
						}
						else{
							ans[now] = data[2].second;
						}
						visited[now] = true;
						for(auto next: edge[now]){
							if(visited[next]) continue;
							if(ans[next] != data[i].second) continue;
							modify(next);
						}
					};
					modify(j);
					break;
				}
			}
		}
		return;
	};

	vi subtree(N);

	std::function<void(ll)> dfs_subtree = [&](ll now){
		subtree[now] = 1;
		for(auto next: edge[now]){
			if(subtree[next]) continue;
			dfs_subtree(next);
			subtree[now] += subtree[next];
		}
	};

	dfs_subtree(0);

	std::function<void(ll, ll)> dfs_coloring = [&](ll now, ll color){
		ans[now] = color;
		for(auto next: edge[now]){
			if(subtree[next] > subtree[now]) continue;
			dfs_coloring(next, color);
		}
	};

	rep(i,1,N){
		if(subtree[i] < data[0].first) continue;
		if(N-subtree[i] < data[1].first) continue;
		ans.assign(N,data[1].second);
		dfs_coloring(i, data[0].second);
		modify_all();
		return ans;
	}

	rep(i,1,N){
		if(subtree[i] < data[1].first) continue;
		if(N-subtree[i] < data[0].first) continue;
		ans.assign(N, data[0].second);
		dfs_coloring(i, data[1].second);
		modify_all();
		return ans;
	}

    rep(i,0,N){
        bool flag = true;
        for(auto next: edge[i]){
            if(subtree[next]*2 > N) flag = false;
        }
        if((N-subtree[i])*2 > N) flag = false;
        if(!flag) continue;
        // centroid
        union_find tree2(N);
        rep(j,0,N){
            for(auto next: edge[j]){
                if(j == i) continue;
                if(next == i) continue;
                tree2.merge(j,next);
            }
        }
        rep(j,0,u.size()){
            if(u[j] == i) continue;
            if(v[j] == i) continue;
            tree2.merge(u[j], v[j]);
            if(tree2.size(u[j]) >= data[0].first && N-tree2.size(u[j]) >= data[1].first){
                //assert(N-tree2.size(u[j]) >= data[1].first);
                rep(k,0,u.size()){
                    edge[u[k]].pb(v[k]);
                    edge[v[k]].pb(u[k]);
                }
                rep(k,0,N){
                    if(tree2.same(u[j], k)) ans[k] = data[0].second;
                    else ans[k] = data[1].second;
                }
                modify_all();
                return ans;
            }
            if(tree2.size(u[j]) >= data[1].first && N-tree2.size(u[j]) >= data[0].first){
                //assert(N-tree2.size(u[j]) >= data[1].first);
                rep(k,0,u.size()){
                    edge[u[k]].pb(v[k]);
                    edge[v[k]].pb(u[k]);
                }
                rep(k,0,N){
                    if(tree2.same(u[j], k)) ans[k] = data[1].second;
                    else ans[k] = data[0].second;
                }
                modify_all();
                return ans;
            }
        }
    }

	ans.assign(N,0);
	return ans;
}

• Subtask 1: 7 / 7
• Subtask 2: 11 / 11
• Subtask 3: 22 / 22
• Subtask 4: 0 / 24
• Subtask 5: 0 / 36