Submission #656562

#	Time	Username	Problem	Language	Result	Execution time	Memory
656562		jiahng	Split the Attractions (IOI19_split)	C++14	100 / 100	145 ms	33612 KiB

split

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

typedef long long ll;
typedef pair<int,int> pi;
typedef vector <int> vi;
typedef vector <pi> vpi;
typedef pair<pi, ll> pii;
typedef set <ll> si;
typedef long double ld;
#define f first
#define s second
//#define mp make_pair
#define FOR(i,s,e) for(int i=s;i<=int(e);++i)
#define DEC(i,s,e) for(int i=s;i>=int(e);--i)
#define pb push_back
#define all(x) (x).begin(), (x).end()
#define lbd(x, y) lower_bound(all(x), y)
#define ubd(x, y) upper_bound(all(x), y)
#define aFOR(i,x) for (auto i: x)
#define mem(x,i) memset(x,i,sizeof x)
#define fast ios_base::sync_with_stdio(false),cin.tie(0),cout.tie(0)
#define maxn 200010
#define INF (ll)1e9
#define MOD 1000000007
typedef pair <vi, int> pvi;
typedef pair <int,pi> ipi;
typedef vector <pii> vpii;

int N;
vi adj[maxn],adj2[maxn],tree[maxn];
bool vis[maxn];
int sz[maxn];
void dfs(int x,int p){
    vis[x] = 1;
    sz[x] = 1;
    aFOR(i, adj[x]) if (i != p && !vis[i]){
        dfs(i,x);
        sz[x] += sz[i];
        tree[x].pb(i);
        tree[i].pb(x);
    }
}
int get_cent(int x,int p){
    aFOR(i, tree[x]) if (i != p){
        if (sz[i] > N / 2) return get_cent(i, x);
    }
    return x;
}
vi ans;
void assign(int x,int p,int v){
    ans[x] = v;
    aFOR(i, tree[x]) if (i != p) assign(i, x, v);
}
int grp[maxn];
void dfs2(int x,int p,int v){
    grp[x] = v;
    aFOR(i, tree[x]) if (i != p) dfs2(i,x,v);
}
int dfs3(int x){
    vis[x] = 1;
    int ans = sz[x];
    aFOR(i, adj2[x]) if (!vis[i]){
        ans += dfs3(i);
    }
    return ans;
}

int cur, cent;
void mark(int x){
    vis[x] = 1;
    cur -= sz[x];
    assert(ans[x] == 2);
    assign(x, cent, 1);

    if (cur <= 0) return;
    aFOR(i, adj2[x]) if (!vis[i] && cur > 0) mark(i);
}

void dfs4(int x,int y, int v){
    if (cur < 0) return;
    ans[x] = v;
    cur--;
    if (cur == 0) return;
    aFOR(i,adj[x]) if (ans[i] == y && cur > 0) dfs4(i,y,v);
}
void get_sizes(int x,int p){
    sz[x] = 1;
    aFOR(i, tree[x]) if (i != p){
        get_sizes(i, x); sz[x] += sz[i];
    }
}

int m[3];
vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
    N = n;
    ans = vi(N,2);
    FOR(i,0,p.size()-1){
        adj[p[i]].pb(q[i]); adj[q[i]].pb(p[i]);
    }
    FOR(i,0,N-1) vis[i] = 0;
    dfs(0,-1); // Get sizes rooted at 0 and spanning tree (Stored in tree array)
    cent = get_cent(0, -1);  // Get the centroid
    
    vpi v;
    v.pb(pi(a,1)); v.pb(pi(b, 2)); v.pb(pi(c,3));
    sort(all(v));

    get_sizes(cent, -1); // Change sz array to reflect sizes rooted at centroid
    bool done = 0;
    aFOR(i, tree[cent]){
        assert(sz[i] <= N / 2);
        if (sz[i] >= v[0].f){
            assign(i, cent, 1); // If A can fit into a single component, set the entire thing to 1 first
            //assert(sz[i] >= count(all(ans), 1));
            done = 1; break;
        }
    }
    if (!done){
        aFOR(i, tree[cent]){
            dfs2(i,cent,i); // Assign groups to each node
        }
        FOR(i,0,p.size()-1) if (p[i] != cent && q[i] != cent && grp[p[i]] != grp[q[i]]){ // Build graph between centroid sons
            adj2[grp[p[i]]].pb(grp[q[i]]);
            adj2[grp[q[i]]].pb(grp[p[i]]);
        } 
        FOR(i,0,N-1) vis[i] = 0;
        aFOR(i, tree[cent]){
            int x = dfs3(i); // Get the total weight of the component under adj2
            if (x >= v[0].f){
                FOR(j,0,N-1) vis[j] = 0;
                cur = v[0].f;
                
                mark(i);
                //assert(N - (v[0].f - cur) >= v[1].f);
                //assert((v[0].f - cur) == count(all(ans), 1)); 
                done = 1;
                break;
            }
        }
    }
    if (!done) return vi(N,0);
    //assert(count(all(ans), 1) >= v[0].f);
    //assert(count(all(ans), 2) >= v[1].f);
 
    FOR(i,1,2){
        cur = v[i-1].f; if (cur == 0) continue;
        int x = find(all(ans), i) - ans.begin();
        dfs4(x,i,i+2);
    }
    FOR(i,0,N-1){
        if (ans[i] < 3) ans[i] = 3;
        else ans[i] -= 2;

        ans[i] = v[ans[i]-1].s;
    }

    assert(count(all(ans), 1) == a);
    assert(count(all(ans), 2) == b);
    assert(count(all(ans), 3) == c);

    return ans;
   
    
    
    
}

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