Submission #1233231

#	Time	Username	Problem	Language	Result	Execution time	Memory
1233231		tapilyoca	Split the Attractions (IOI19_split)	C++20	100 / 100	148 ms	30136 KiB

split

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

template<typename T>
using vec = vector<T>;
using ll = long long;
using vll = vec<ll>;
using vvll = vec<vll>;
using pll = pair<ll,ll>;
#define pb push_back
#define dbg(x) cerr << #x << ": " << x << endl
// #define cerr if(0) cout

/* type shit */

ll curr = 10, N, centroid = -1, centroid2 = -1;
vll subtree_size, vis, parent, dayOrder(3);
vec<int> color;
vec<bool> moveable, moveable2, moved, is_in_centroid_subtree, is_in_centroid2_subtree;
vec<pll> days;
vvll adj, dt_adj;
map<ll,bool> colorworks;

pll c1 = {1e18,-1}, c2 = {1e18,-1}; // ST_SIZE, NODE INDEX

pll yeah(pll a, pll b) {
    if(a.first < b.first) return a;
    return b;
}

void dfs(ll u, ll p) {
    parent[u] = p;

    for(ll &v : adj[u]) {
        if(v == p or vis[v]) continue;
        vis[v] = 1;
        dt_adj[u].pb(v);
        dt_adj[v].pb(u);
        dfs(v,u);
    }
}

void dfs_tree(ll u, ll p) {
    subtree_size[u] = 1;

    for(ll &v : dt_adj[u]) {
        if(v == p) continue;
        dfs_tree(v,u);
        subtree_size[u] += subtree_size[v];
    }
    
    if(centroid == -1 and subtree_size[u] >= days[0].first) c1 = yeah(c1, {subtree_size[u], u});

    if(centroid2 == -1 and subtree_size[u] >= days[1].first) c2 = yeah(c2, {subtree_size[u], u});
}

void dfs_cassign(ll u, ll p) {
    is_in_centroid_subtree[u] = is_in_centroid_subtree[p];
    is_in_centroid2_subtree[u] = is_in_centroid2_subtree[p];
    
    if(u == centroid) is_in_centroid_subtree[u] = 1;
    if(u == centroid2) is_in_centroid2_subtree[u] = 1;
    for(ll &v : dt_adj[u]) {
        if(v == p) continue;
        dfs_cassign(v,u);
    }
}

bool check_moveable(ll x, ll ct) {
    queue<ll> q;

    q.push(x);
    vis[x] = curr;
    while(!q.empty()) {
        ll u = q.front();
        q.pop();

        for(ll &v : adj[u]) {
            if(vis[v] or v == ct) continue;
            vis[v] = curr;
            if(ct == centroid and !is_in_centroid_subtree[v]) colorworks[curr] = 1;
            else if(ct == centroid2 and !is_in_centroid2_subtree[v]) colorworks[curr] = 1;
            q.push(v);
        }
    }

    return colorworks[curr];
}

void dfs_color_first(ll u, ll p, ll c) {
    if(days[dayOrder[c-1]].first == 0) return;

    color[u] = c;
    days[dayOrder[c-1]].first--;
    for(ll &v : dt_adj[u]) {
        if(v == p or moved[v]) continue;
        dfs_color_first(v,u,c);
    }
}

void bfs_color_second(ll x, ll c) {
    queue<ll> q;
    q.push(x);

    color[x] = c;
    days[dayOrder[c-1]].first--;
    if(days[dayOrder[c-1]].first == 0) return;

    while(!q.empty()) {
        ll u = q.front();
        q.pop();

        for(ll &v : adj[u]) {
            if(color[v] != -1) continue;
            color[v] = c;
            days[dayOrder[c-1]].first--;
            if(days[dayOrder[c-1]].first == 0) return;
            q.push(v);
        }
    }
}

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
	  N = n;
    ll m = p.size();
    adj.resize(n);
    dt_adj.resize(n);
    parent.assign(n,-1);
    vis.assign(n,0);
    moveable.assign(n,0);
    moveable2.assign(n,0);
    moved.assign(n,0);
    is_in_centroid_subtree.assign(n,0);
    is_in_centroid2_subtree.assign(n,0);

    for(int i = 0; i < m; i++) {
        adj[p[i]].pb(q[i]);
        adj[q[i]].pb(p[i]);
    }

    days = {{a,1}, {b,2}, {c,3}};
    sort(days.begin(),days.end());
    for(int i = 0; i < 3; i++) dayOrder[days[i].second - 1] = i;

    subtree_size.assign(n,0);
    color.assign(n,-1);
    vis[0] = 1;

    dfs(0,-1);
    dfs_tree(0,-1);

    centroid = c1.second;
    centroid2 = c2.second;

    if(centroid == -1 and centroid2 == -1) {
        vec<int> ans(n,0);
        return ans;
    }

    dfs_cassign(0,0);
    vis.assign(n,0);

    if(centroid != -1) {
        for(ll &x : dt_adj[centroid]) {
            if(x == parent[centroid]) continue;
            if(vis[x]) moveable[x] = colorworks[vis[x]];
            else moveable[x] = check_moveable(x, centroid);
            curr++;
        }
    }

    vis.assign(n,0);

    if(centroid2 != -1) {
        for(ll &x : dt_adj[centroid2]) {
            if(x == parent[centroid2] or vis[x]) continue;
            if(vis[x]) moveable2[x] = colorworks[vis[x]];
            else moveable2[x] = check_moveable(x, centroid2);
            curr++;
        }
    }

    // CASE: A fits in centroid subtree, B fits above centroid
    ll c_st_size = -1;
    ll rest = -1;
    if(centroid != -1) {
        c_st_size = subtree_size[centroid];
        rest = n - c_st_size;
        for(ll &x : dt_adj[centroid]) {
            if(rest >= days[1].first and c_st_size >= days[0].first) break;
            if(x == parent[centroid] or not moveable[x]) continue;
            if(c_st_size - subtree_size[x] >= days[0].first) {
                moved[x] = 1; 
                c_st_size -= subtree_size[x];
                rest += subtree_size[x];
            }
        }

        if(c_st_size >= days[0].first and rest >= days[1].first) {
            dfs_color_first(centroid, parent[centroid], days[0].second);
            bfs_color_second(parent[centroid], days[1].second);
            for(int &x : color) x = (x != -1 ? x : days[2].second);
            return color;
        }
    }

    // CASE: B fits in centroid subtree, A fits above centroid
    if(centroid2 != -1) {
        c_st_size = subtree_size[centroid2];
        rest = n - c_st_size;
        moved.assign(n,0);

        for(ll &x : dt_adj[centroid2]) {
            if(rest >= days[0].first and c_st_size >= days[1].first) break;
            if(x == parent[centroid2] or not moveable2[x]) continue;
            if(c_st_size - subtree_size[x] >= days[1].first) {
                moved[x] = 1;                
                c_st_size -= subtree_size[x];
                rest += subtree_size[x];
            }
        }

        if(c_st_size >= days[1].first and rest >= days[0].first) {
            dfs_color_first(centroid2,parent[centroid2],days[1].second);
            bfs_color_second(parent[centroid2],days[0].second);
            for(int &x : color) x = (x != -1 ? x : days[2].second);
            return color;
        }
    }

    color.assign(n,0);
    return color;
}

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