Submission #386052

#	Time	Username	Problem	Language	Result	Execution time	Memory
386052		rocks03	Split the Attractions (IOI19_split)	C++14	40 / 100	744 ms	1048580 KiB

split

#include<bits/stdc++.h>
using namespace std;
#define ll long long
#define pii pair<int, int>
#define pll pair<ll, ll>
#define ff first
#define ss second
#define pb push_back
#define SZ(x) ((int)(x).size())
#define all(x) x.begin(), x.end()
#define debug(x) cout << #x << ": " << x << " "
#define nl cout << "\n"
#define rep(i, a, b) for(int i = (a); i < (b); i++)
#define per(i, a, b) for(int i = (a); i >= (b); i--)
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

const int MAXN = 2e5+100;
vector<int> g[MAXN];
vector<int> ans;

int sz[MAXN];
void getsz(int v, int p){
    sz[v] = 1;
    for(int u : g[v]){
        if(u ^ p){
            getsz(u, v);
            sz[v] += sz[u];
        }
    }
}

bool done = false;
int N, a, b, c;

void dfs(int v, int p, int col, int& x){
    if(!x) return;
    ans[v] = col, --x;
    for(int u : g[v]){
        if(u ^ p){
            if(!ans[u]){
                dfs(u, v, col, x);
            }
        }
    }
}

void solve(int v, int p){
    for(int u : g[v]){
        if(u ^ p){
            if(sz[u] >= a){
                if(N - sz[u] >= b){
                    dfs(u, v, 1, a);
                    dfs(v, -1, 2, b);
                    rep(i, 0, N){
                        if(!ans[i]) ans[i] = 3;
                    }
                    done = true; return;
                } else if(N - sz[u] >= c){
                    dfs(u, v, 1, a);
                    dfs(v, -1, 3, c);
                    rep(i, 0, N){
                        if(!ans[i]) ans[i] = 2;
                    }
                    done = true; return;
                }
            }
            if(sz[u] >= b){
                if(N - sz[u] >= a){
                    dfs(u, v, 2, b);
                    dfs(v, -1, 1, a);
                    rep(i, 0, N){
                        if(!ans[i]) ans[i] = 3;
                    }
                    done = true; return;
                } else if(N - sz[u] >= c){
                    dfs(u, v, 2, b);
                    dfs(v, -1, 3, c);
                    rep(i, 0, N){
                        if(!ans[i]) ans[i] = 1;
                    }
                    done = true; return;
                }
            }
            if(sz[u] >= c){
                if(N - sz[u] >= b){
                    dfs(u, v, 3, c);
                    dfs(v, -1, 2, b);
                    rep(i, 0, N){
                        if(!ans[i]) ans[i] = 1;
                    }
                    done = true; return;
                } else if(N - sz[u] >= a){
                    dfs(u, v, 3, c);
                    dfs(v, -1, 1, a);
                    rep(i, 0, N){
                        if(!ans[i]) ans[i] = 2;
                    }
                    done = true; return;
                }
            }
        }
    }
    for(int u : g[v]){
        if(u ^ p){
            int oldsz = sz[v];
            sz[v] = N - sz[u];
            solve(u, v);
            sz[v] = oldsz;
            if(done) return;
        }
    }
}

vector<int> find_split(int N, int a, int b, int c, vector<int> p, vector<int> q) {
    rep(i, 0, SZ(p)){
        g[p[i]].pb(q[i]);
        g[q[i]].pb(p[i]);
    }
    ans = vector<int>(N, 0);
    bool subtask1 = true;
    rep(i, 0, N) subtask1 &= (SZ(g[i]) <= 2);
    if(subtask1){
        int v = 0;
        rep(i, 0, N){
            if(SZ(g[i]) == 1) v = i;
        }
        queue<int> q; q.push(v);
        while(!q.empty()){
            int v = q.front();
            q.pop();
            if(a) ans[v] = 1, --a;
            else if(b) ans[v] = 2, --b;
            else if(c) ans[v] = 3, --c;
            for(int u : g[v]){
                if(!ans[u]) q.push(u);
            }
        }
        return ans;
    }
    if(a == 1){
        queue<int> q; q.push(0);
        ans[0] = 2, --b;
        while(!q.empty()){
            int v = q.front();
            q.pop();
            for(int u : g[v]){
                if(!ans[u]){
                    if(b) ans[u] = 2, --b;
                    else if(c) ans[u] = 3, --c;
                    else if(a) ans[u] = 1, --a;
                    q.push(u);
                }
            }
        }
        return ans;
    }
    if(SZ(p) == N - 1){
        ::N = N, ::a = a, ::b = b, ::c = c;
        getsz(0, -1);
        solve(0, -1);
        return ans;
    }
}

Compilation message (stderr)

split.cpp: In function 'std::vector<int> find_split(int, int, int, int, std::vector<int>, std::vector<int>)':
split.cpp:163:1: warning: control reaches end of non-void function [-Wreturn-type]
  163 | }
      | ^

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