Submission #297570

#	Time	Username	Problem	Language	Result	Execution time	Memory
297570		shayan_p	Split the Attractions (IOI19_split)	C++17	100 / 100	188 ms	17400 KiB

split

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

#define F first
#define S second
#define PB push_back
#define sz(s) (int(s.size()))
#define bit(n, k) (((n)>>(k))&1)

using namespace std;

typedef long long ll;
typedef pair<int, int> pii;

const int maxn = 1e5 + 10;

struct DSU{
    int pr[maxn];
    DSU(){
	memset(pr, -1, sizeof pr);
    }
    int fnd(int u){
	return pr[u] < 0 ? u : pr[u] = fnd(pr[u]);
    }
    void mrg(int a, int b){
	if( (a = fnd(a)) == (b = fnd(b)) )
	    return;
	if(pr[a] > pr[b])
	    swap(a, b);
	pr[a]+= pr[b];
	pr[b] = a;
    }
    int size(int u){
	return -pr[fnd(u)];
    }
};
DSU dsu;

vector<int> v[maxn];

int pr[maxn], SZ[maxn];
bool mark[maxn];
int col[maxn];

void dfs(int u){
    SZ[u] = 1;
    mark[u] = 1;
    for(int y : v[u]){
	if(!mark[y])
	    pr[y] = u, dfs(y), SZ[u]+= SZ[y];
    }
}

int color(int u, int cnt, int c){
    assert(cnt);
    cnt--;
    col[u] = c;
    for(int y : v[u]){
	if(mark[y] == mark[u] && col[y] != c && cnt)
	    cnt = color(y, cnt, c);
    }
    return cnt;
}

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q){
    vector<pii> SZS;
    SZS.PB({a, 1});
    SZS.PB({b, 2});
    SZS.PB({c, 3});
    sort(SZS.begin(), SZS.end());

    int A = SZS[0].F;

    auto mkay = [&](){
		    int cnt = 0, ZR = -1, ON = -1;
		    for(int i = 0; i < n; i++){
			if(mark[i] == 0)
			    ZR = i;
			else
			    ON = i;
			cnt+= mark[i];
		    }
		    assert(A <= cnt && cnt <= n-A);
		    if(cnt > SZS[1].F)
			swap(ZR, ON);
		    color(ON, SZS[0].F, SZS[0].S);
		    color(ZR, SZS[1].F, SZS[1].S);
		    vector<int> ans;
		    for(int i = 0; i < n; i++){
			if(col[i] == 0)
			    col[i] = SZS[2].S;
			ans.PB(col[i]);
		    }
		    return ans;
		};

    int m = sz(p);
    for(int i = 0; i < m; i++){
	v[p[i]].PB(q[i]);
	v[q[i]].PB(p[i]);
    }

    dfs(0);
    for(int u = 0; u < n; u++){
	if(A <= SZ[u] && SZ[u] <= n-A){
	    for(int i = 0; i < n; i++)
		v[i].clear();
	    for(int i = 1; i < n; i++)
		v[i].PB(pr[i]), v[pr[i]].PB(i);
	    memset(mark, 0, sizeof mark);
	    mark[pr[u]] = 1;
	    dfs(u);
	    mark[pr[u]] = 0;
	    return mkay();
	}
    }

    int root = 0;
    for(int i = 0; i < n; i++){
	if(SZ[i] > n-A && SZ[root] > SZ[i])
	    root = i;
    }

    for(int i = 1; i < n; i++){
	if(i != root && pr[i] != root){
	    dsu.mrg(i, pr[i]);	    
	}
    }
    
    for(int i = 0; i < m; i++){
	if(p[i] == root || q[i] == root)
	    continue;
	dsu.mrg(p[i], q[i]);
	if(dsu.size(p[i]) >= A){
	    memset(mark, 0, sizeof mark);
	    for(int j = 0; j < n; j++){
		if(dsu.fnd(j) == dsu.fnd(p[i]))
		    mark[j] = 1;
	    }		
	    return mkay();
	}
    }

    vector<int> ans(n);
    return ans;
}

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