Submission #289909

#	Time	Username	Problem	Language	Result	Execution time	Memory
289909		MarcoMeijer	Split the Attractions (IOI19_split)	C++14	100 / 100	218 ms	21880 KiB

split

#include <bits/stdc++.h>
#include "split.h"
using namespace std;

//macros
typedef long long ll;
typedef pair<int, int> ii;
typedef pair<ll, ll> lll;
typedef tuple<int, int, int> iii;
typedef vector<int> vi;
typedef vector<ii> vii;
typedef vector<iii> viii;
typedef vector<ll> vll;
typedef vector<lll> vlll;
#define REP(a,b,c) for(int a=int(b); a<int(c); a++)
#define RE(a,c) REP(a,0,c)
#define RE1(a,c) REP(a,1,c+1)
#define REI(a,b,c) REP(a,b,c+1)
#define REV(a,b,c) for(int a=int(c-1); a>=int(b); a--)
#define FOR(a,b) for(auto& a:b)
#define INF 1e9
#define pb push_back
#define fi first
#define se second

const int MX = 1e5+10;

int n, m, a[3], SA[3];
int cnt[3];
int RSA[4];
vi adj[MX];
vi adjT[MX];
vi res;
bitset<MX> visited;
int sz[MX];
int cent;

void dfsTree(int u) {
    visited[u] = 1;
    FOR(v,adj[u]) if(!visited[v]) dfsTree(v), adjT[u].pb(v), adjT[v].pb(u);
}
void dfsSz(int u, int p) {
    sz[u] = 1;
    FOR(v,adjT[u]) if(v != p) dfsSz(v, u), sz[u]+=sz[v];
}
int dfsCentroid(int u, int p, int s) {
    FOR(v,adjT[u]) if(v != p && sz[v] >= s/2) return dfsCentroid(v,u,s);
    return u;
}

int comp[MX], compSZ[MX], compR[MX], compP[MX], compN=0;
void dfsComp(int u, int p) {
    comp[u] = compN;
    FOR(v,adjT[u]) if(v != p) dfsComp(v,u);
}
int getComp(int i) {return i==compP[i] ? i : compP[i] = getComp(compP[i]);}
bool isSameSet(int i, int j) {return getComp(i) == getComp(j);}
void unionSet(int i, int j) {
    if(isSameSet(i,j)) return;
    i=compP[i], j=compP[j];
    compSZ[i] = compSZ[j] = compSZ[i] + compSZ[j];
    if(compR[i] > compR[j]) {
        compP[j] = i;
    } else {
        compP[i] = j;
        if(compR[i] == compR[j]) compR[j]++;
    }
}

void dfsFirst(int u, int c, int numb) {
    if(res[u] != 0) return;
    if(cnt[numb] == a[SA[numb]]) return;
    cnt[numb]++;
    res[u] = numb+1;
    FOR(v,adj[u]) if(v != cent && getComp(comp[v]) == c) dfsFirst(v,c,numb);
}
void dfsSecond(int u, int numb) {
    if(res[u] != 0) return;
    if(cnt[numb] == a[SA[numb]]) return;
    cnt[numb]++;
    res[u] = numb+1;
    FOR(v,adj[u]) dfsSecond(v, numb);
}

vi find_split(int N, int A, int B, int C, vi P, vi Q) {
    n=N; a[0]=A; a[1]=B; a[2]=C;
    RE(i,3) SA[i]=i;
    sort(SA,SA+3,[](int i, int j) {
        return a[i]<a[j];
    });
    res.assign(n,0);
    m=P.size();
    RE(i,m) {
        adj[P[i]].pb(Q[i]);
        adj[Q[i]].pb(P[i]);
    }
    visited.reset(); dfsTree(1);
    dfsSz(1, -1);
    cent = dfsCentroid(1, -1, sz[1]);
    dfsSz(cent, -1);
    FOR(u,adjT[cent]) {
        compSZ[compN] = sz[u];
        dfsComp(u,cent);
        compN++;
    }
    RE(i,compN) compP[i]=i, compR[i]=1;
    RE(i,m) {
        if(P[i] == cent || Q[i] == cent) continue;
        int u = getComp(comp[P[i]]);
        int v = getComp(comp[Q[i]]);
        if(compSZ[u] < a[SA[0]] && compSZ[v] < a[SA[0]])
            unionSet(u,v);
    }

    int firstComp = -1;
    RE(i,compN)
        if(compSZ[i] >= a[SA[0]])
            firstComp = i;
    if(firstComp == -1) return res;

    RE(i,3) cnt[i] = 0;
    RE(i,n) if(i != cent && compSZ[getComp(comp[i])] >= a[SA[0]]) {
        dfsFirst(i, getComp(comp[i]),0);
        break;
    }
    dfsSecond(cent,1);
    RE(i,n) if(res[i] == 0) res[i] = 3;
    RE(i,n) res[i] = SA[res[i]-1]+1;

	return res;
}

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