Submission #970746

#	Time	Username	Problem	Language	Result	Execution time	Memory
970746		GrindMachine	Split the Attractions (IOI19_split)	C++17	100 / 100	83 ms	29508 KiB

split

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace std;
using namespace __gnu_pbds;

template<typename T> using Tree = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;
typedef long long int ll;
typedef long double ld;
typedef pair<int,int> pii;
typedef pair<ll,ll> pll;

#define fastio ios_base::sync_with_stdio(false); cin.tie(NULL)
#define pb push_back
#define endl '\n'
#define sz(a) (int)a.size()
#define setbits(x) __builtin_popcountll(x)
#define ff first
#define ss second
#define conts continue
#define ceil2(x,y) ((x+y-1)/(y))
#define all(a) a.begin(), a.end()
#define rall(a) a.rbegin(), a.rend()
#define yes cout << "Yes" << endl
#define no cout << "No" << endl

#define rep(i,n) for(int i = 0; i < n; ++i)
#define rep1(i,n) for(int i = 1; i <= n; ++i)
#define rev(i,s,e) for(int i = s; i >= e; --i)
#define trav(i,a) for(auto &i : a)

template<typename T>
void amin(T &a, T b) {
    a = min(a,b);
}

template<typename T>
void amax(T &a, T b) {
    a = max(a,b);
}

#ifdef LOCAL
#include "debug.h"
#else
#define debug(x) 42
#endif

/*

refs:
https://codeforces.com/blog/entry/68940?#comment-532790
https://codeforces.com/blog/entry/68940?#comment-532806

*/

const int MOD = 1e9 + 7;
const int N = 1e5 + 5;
const int inf1 = int(1e9) + 5;
const ll inf2 = ll(1e18) + 5;

#include "split.h"

vector<int> adj1[N], adj2[N], adj3[N];
vector<bool> vis(N);

void dfs1(int u){
    vis[u] = 1;
    trav(v,adj1[u]){
        if(vis[v]) conts;
        adj2[u].pb(v), adj2[v].pb(u);
        dfs1(v);
    }
}

vector<int> subsiz(N);

void dfs2(int u, int p){
    subsiz[u] = 1;
    trav(v,adj2[u]){
        if(v == p) conts;
        dfs2(v,u);
        subsiz[u] += subsiz[v];
    }
}

int dfs3(int u, int p){
    trav(v,adj2[u]){
        if(v == p) conts;
        if(subsiz[v] > subsiz[0]/2){
            return dfs3(v,u);
        }
    }

    return u;
}

vector<int> nodes[N];
vector<int> subr(N);

void dfs4(int u, int p, int r){
    nodes[r].pb(u);
    subr[u] = r;
    trav(v,adj2[u]){
        if(v == p) conts;
        dfs4(v,u,r);
    }
}

vector<int> ans;
int to_col;

void dfs5(int u, int p){
    if(to_col){
        ans[u] = 1;
        to_col--;
    }

    trav(v,adj2[u]){
        if(v == p) conts;
        if(ans[v] != -1) conts;
        dfs5(v,u);
    }
}

vector<int> ord;

void dfs6(int u){
    vis[u] = 1;
    ord.pb(u);
    trav(v,adj3[u]){
        if(vis[v]) conts;
        dfs6(v);
    }
}

vector<int> find_split(int n, int A, int B, int C, vector<int> U, vector<int> V) {
    rep(i,sz(U)){
        int u = U[i], v = V[i];
        adj1[u].pb(v), adj1[v].pb(u);
    }

    // find arbitrary spanning tree
    dfs1(0);

    // find centroid of spanning tree
    dfs2(0,-1);
    int c = dfs3(0,-1);
        
    // root @centroid
    // wlog, A <= B <= C
    // if any of the subtrees have size >= A, done
    // we can find a cc of size A in that subtree and this may waste at most n/2 nodes (cuz rooted @centroid)
    // B <= n/2, so we can find a cc of size B in the rest of the tree
    vector<int> subs;
    trav(v,adj2[c]){
        subs.pb(v);
        dfs4(v,c,v);
    }

    vector<pii> col_ord = {{A,0},{B,1},{C,2}};
    sort(all(col_ord));
    A = col_ord[0].ff, B = col_ord[1].ff, C = col_ord[2].ff;
    ans = vector<int>(n,-1);

    trav(v,subs){
        if(sz(nodes[v]) < A) conts;
        rep(j,A){
            ans[nodes[v][j]] = 0;
        }

        to_col = B;
        dfs5(c,-1);
        rep(i,n) if(ans[i] == -1) ans[i] = 2;

        rep(i,n) ans[i] = col_ord[ans[i]].ss+1;
        return ans;
    }

    // if no solution found, it means that all subtrees have size < A
    // consider edges that go between subtrees
    // each subtree compressed into a single node with weight = size of subtree
    // add edge between 2 nodes if the corresponding subtrees are connected by an edge
    // find a connected subgraph of the nodes with sum of weights >= A
    // found using a dfs, stop at the first time when size >= A
    // at most 2A nodes wasted after this (because all have weight < A)
    // A+B+C = n, A <= C, so 2A+B <= n
    // so we will be able to find a connected subgraph of size B in the remaining graph
    // if no connected subgraph has sum of weights >= A, then no solution exists
    rep(u,n){
        trav(v,adj1[u]){
            if(u == c or v == c) conts;
            if(subr[u] != subr[v]){
                adj3[subr[u]].pb(subr[v]);
            }
        }
    }

    fill(all(vis),0);

    trav(u,subs){
        if(vis[u]) conts;
        ord.clear();
        dfs6(u);
        int sum = 0;
        trav(v,ord) sum += sz(nodes[v]);
        if(sum < A) conts;
        sum = 0;

        trav(v,ord){
            auto &curr = nodes[v];
            rep(j,sz(curr)){
                if(sum == A) break;
                ans[curr[j]] = 0;
                sum++;
            }
        }

        assert(sum == A);

        to_col = B;
        dfs5(c,-1);
        rep(i,n) if(ans[i] == -1) ans[i] = 2;

        rep(i,n) ans[i] = col_ord[ans[i]].ss+1;
        return ans;
    }

    fill(all(ans),0);
    return ans;
}

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