Submission #1233838

#	Time	Username	Problem	Language	Result	Execution time	Memory
1233838		jasonic	Split the Attractions (IOI19_split)	C++20	40 / 100	94 ms	20296 KiB

split

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

#define ll long long
#define fastIO cin.tie(0); ios::sync_with_stdio(false)

int labels[3] = {1, 2, 3};
int n, m, a, b, c;

vector<vector<int>> bigadjlist;
vector<vector<int>> adjlist;
int sz[200005];

void calc_sizes(int v, int p = -1) {
    sz[v] = 1;

    for(auto i : adjlist[v]) if(i != p) {
        calc_sizes(i, v);
        sz[v] += sz[i];
    }
};

int get_centroid(int v, int p = -1) {
    for(auto i : adjlist[v]) if(i != p) {
        if(sz[i] * 2 > n) return get_centroid(i, v);
    }
    return v;
}

int need = 0;
void dfs_fill(vector<int> &arr, int val, int v, int p = -1) {
    if(need == 0) return;
    arr[v] = val;
    need--;
    if(need == 0) return;

    for(auto i : adjlist[v]) if(i != p) {
        dfs_fill(arr, val, i, v);
        if(need == 0) return;
    }
}

vector<int> solve(int c, int s) {
    vector<int> ans(n, labels[2]);
    vector<bool> vis(n, false);
    need = a;
    dfs_fill(ans, labels[0], s, c);
    need = b;
    dfs_fill(ans, labels[1], c, s);
    return ans;
}

vector<int> find_split(int N, int A, int B, int C, vector<int> p, vector<int> q) {
    n = N;
    m = p.size();
    a=A, b=B, c=C;

    bigadjlist = vector<vector<int>>(n);
    adjlist = vector<vector<int>>(n);

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

    // assign the correct things
    if(c < a) {
        swap(a, c);
        swap(labels[0], labels[2]);
    }

    if(b < a) {
        swap(a, b);
        swap(labels[0], labels[1]);
    }

    if(c < b) {
        swap(c, b);
        swap(labels[1], labels[2]);
    }

    // dfs tree to find correct edges
    vector<int> vis(n, false);
    stack<int> dfs;
    dfs.push(0);
    vis[0] = true;
    while(!dfs.empty()) {
        int curr = dfs.top(); dfs.pop();
        for(auto i : bigadjlist[curr]) if(!vis[i]) {
            adjlist[curr].push_back(i);
            adjlist[i].push_back(curr);
            dfs.push(i);
            vis[i] = true;
        }
    }

    calc_sizes(0);
    int centroid = get_centroid(0);
    calc_sizes(centroid);

    // check if the dfs tree is of any use
    for(auto i : adjlist[centroid]) {
        if(a <= sz[i]) return solve(centroid, i);
    }

    vector<int> ans = vector<int>(n, 0);

    return ans;
}

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