제출 #1190772

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
1190772		Swan	Split the Attractions (IOI19_split)	C++20	100 / 100	89 ms	25280 KiB

split

#pragma GCC optimize("O3")
#pragma GCC optimize("unroll-loops")

#include <iostream>
#include <vector>
#include <stack>
#include <algorithm>
#include <string>
#include <set>
#include <map>
#include <list>
#include <time.h>
#include <math.h>
#include <random>
#include <deque>
#include <queue>
#include <cassert>
#include <climits>
#include <unordered_map>
#include <unordered_set>
#include <iomanip>
#include <bitset>
#include <sstream>
#include <chrono>
#include <cstring>
#include <optional>
//#define int long long
#define INP freopen("subrev.in","r",stdin)
#define OUTP freopen("subrev.out","w",stdout)
using ld = long double;
using LD = ld;
 
using namespace std;

const int maxn = 2e5 + 123;
int N;
vector<vector<int> > v;
vector<vector<int> > tree;
bitset<maxn> used;
vector<int> colors;
int sz[maxn];
int compColor[maxn];

void dfs0(int u) {
    used[u] = 1;
    sz[u] = 1;

    for (auto a : v[u]) {
        if (used[a]) continue;

        dfs0(a);

        tree[u].push_back(a);
        tree[a].push_back(u);        

        sz[u] += sz[a];
    }
}

void getTree(int root) {
    for (int i = 0; i < N; i++) sz[i] = 0;
    dfs0(root);
}

void dfs1(int u) {
    used[u] = 1;
    sz[u] = 1;

    for (auto a : tree[u]) {
        if (used[a]) continue;

        dfs1(a);  

        sz[u] += sz[a];
    }
}

void recalcTree(int root) {
    for (int i = 0; i < N; i++) sz[i] = 0;
    dfs1(root);
}

int getCentroid(int u, int treeSize, int p = -1) {
    for (auto a : tree[u]) {
        if (a == p) continue;
        if (sz[a] > treeSize / 2) {
            return getCentroid(a, treeSize, u);
        }
    }

    return u;
}   

vector<int> vertices;

void getColorTree(int u, int color, int* sz) {
    if ((*sz) == 0) return;
    used[u] = 1;

    colors[u] = color;
    vertices.push_back(u);

    (*sz)--;    

    for (auto a : tree[u]) {
        if (used[a]) continue;
        getColorTree(a, color, sz);
    }
}

void getColor(int u, int color, int* sz) {
    if ((*sz) == 0) return;
    used[u] = 1;

    colors[u] = color;
    vertices.push_back(u);

    (*sz)--;    

    for (auto a : v[u]) {
        if (used[a]) continue;
        getColor(a, color, sz);
    }
}

void getColor2(int u, int color, int* sz) {

    priority_queue<pair<int, int>> q; 
    q.push({-compColor[u], u});

    while (q.size()) {
        if ((*sz) == 0) return;
        int u = q.top().second;
        q.pop();

        if (used[u]) continue;

        used[u] = 1;
        colors[u] = color;
        (*sz)--;    
        vertices.push_back(u);

        for (auto a : v[u]) {
            if (used[a]) continue;
            q.push({-compColor[a], a});
            
        }
    } 
}

void assignColors(int u, int color, int p = -1) {
    compColor[u] = color;
    for (auto a : tree[u]) {
        if (a == p) continue;
        assignColors(a, color, u);
    }
}

bool validate(vector<int> colors, vector<pair<int, int>> mda) {
    
    int a = 0;
    int b = 0;
    int c = 0;

    for (auto x : colors) {
        if (x == mda[0].second) a++;
        if (x == mda[1].second) b++;
        if (x == mda[2].second) c++;
    }

    return (a == mda[0].first && b == mda[1].first && c == mda[2].first);
}

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
    v.resize(n + 10);
    tree.resize(n + 10);
    colors.resize(n, 0);

    vector<pair<int, int>> mda = {{a, 1}, {b, 2}, {c, 3}}; sort(mda.begin(), mda.end());
    a = mda[0].first;
    b = mda[1].first;
    c = mda[2].first;

    N = n;
    
    int A = a;
    int B = b;

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

    getTree(0);

    int centroid = getCentroid(0, n);

    used.reset();

    recalcTree(centroid);

    used.reset();

    int curColor = 1;


    for (int i = 0; i < tree[centroid].size(); i++) {
        int to = tree[centroid][i];
        if (sz[to] >= a) {
            A = a;
            B = b;
            vertices.clear();

            used[centroid] = 1;
            getColorTree(to, mda[0].second, &A);

            used[centroid] = 0;
            getColorTree(centroid, mda[1].second, &B);

            if (A > 0 || B > 0) continue;
            
            for (int j = 0; j < n; j++) {
                if (colors[j] == 0) colors[j] = mda[2].second;
            }

            if (!validate(colors, mda)) assert(0 == 1);


            return colors;
        }

        assignColors(to, curColor++, centroid);

    }

    used.reset();

    used[centroid] = 1;

    for (int i = 0; i < tree[centroid].size(); i++) {
        int to = tree[centroid][i];
        if (used[to]) continue;

        A = a;
        B = b;
        vertices.clear();

        used[centroid] = 1;
        getColor2(to, mda[0].second, &A);

        if (!A) {


            used[centroid] = 0;
            used.reset();
            for (auto a : vertices) used[a] = 1;
            getColor(centroid, mda[1].second, &B);

            

            for (int j = 0; j < n; j++) {
                if (colors[j] == 0) colors[j] = mda[2].second;
            }

            return colors;
        }

        for (auto a : vertices) colors[a] = 0;
    }



    return colors;
}


/*
void solve() {
    vector<int> res = find_split(9, 4, 2, 3, {0, 0, 0, 0, 0, 0, 1, 3, 4, 5},
        {1, 2, 3, 4, 6, 8, 7, 7, 5, 6});
    
    for (int i = 0; i < res.size(); i++) {
        cout << i << ' ' << res[i] << endl;
    }

    cout << endl;
}

main() {
    ios_base::sync_with_stdio(0);

    solve();
}
     */

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