Submission #385654

#	Time	Username	Problem	Language	Result	Execution time	Memory
385654		jhnah917	Split the Attractions (IOI19_split)	C++14	100 / 100	299 ms	25824 KiB

split

/******************************
Author: jhnah917(Justice_Hui)
g++ -std=c++17 -DLOCAL
******************************/

#ifndef LOCAL
#include "split.h"
#endif

#include <bits/stdc++.h>
#define x first
#define y second
#define all(v) v.begin(), v.end()
#define compress(v) sort(all(v)), v.erase(unique(all(v)), v.end())
#define IDX(v, x) (lower_bound(all(v), x) - v.begin())
using namespace std;

using uint = unsigned;
using ll = long long;
using ull = unsigned long long;
using PII = pair<int, int>;

constexpr int SZ = 101010;

template<size_t _Sz> struct UnionFind {
    int P[_Sz], W[_Sz];
    UnionFind(){ clear(); }
    void clear(){ iota(P, P+_Sz, 0); fill(W, W+_Sz, 1); }
    int find(int v){ return v == P[v] ? v : P[v] = find(P[v]); }
    bool merge(int u, int v){
        u = find(u); v = find(v);
        if(u == v) return false;
        P[u] = v; W[v] += W[u];
        return true;
    }
};

int N, M, A, B, C, S[SZ], chk[SZ];
PII ORD[3];
vector<int> G[SZ], Tree[SZ], Comp[SZ];
UnionFind<SZ> UF;

int get_sz(int v, int b=-1){
    S[v] = 1;
    for(auto i : Tree[v]) if(i != b) S[v] += get_sz(i, v);
    return S[v];
}
int get_cent(int v, int tot, int b=-1){
    for(auto i : Tree[v]) if(i != b && S[i]*2 > tot) return get_cent(i, tot, v);
    return v;
}

vector<int> dfs_subtree(int v, int b){
    vector<int> ret;
    function<void(int,int)> go = [&go,&ret](int v, int b){
        ret.push_back(v);
        for(auto i : Tree[v]) if(i != b) go(i, v);
    };
    go(v, b);
    return ret;
}

void dfs_merge(int v, int b){
    for(auto i : Tree[v]) if(i != b) UF.merge(v, i), dfs_merge(i, v);
}

vector<int> dfs_comp(int v){
    vector<int> ret;
    function<void(int)> go = [&go,&ret](int v){
        chk[v] = 1; ret.push_back(v);
        for(auto i : Comp[v]) if(!chk[i]) go(i);
    };
    go(v);
    return ret;
}

int use[SZ];
vector<int> dfs_graph(int v){
    vector<int> ret;
    function<void(int)> go = [&ret,&go](int v){
        chk[v] = 1; ret.push_back(v);
        for(auto i : G[v]) if(!chk[i] && use[v] == use[i]) go(i);
    };
    go(v);
    return ret;
}

vector<int> get_answer(const vector<int> &V){
    vector<int> res(N, ORD[2].y);
    memset(chk, 0, sizeof chk);

    for(auto i : V) use[i] = 1;
    for(int i=0; i<N; i++){
        if(chk[i]) continue;
        auto dfn = dfs_graph(i);
        dfn.resize(ORD[!use[i]].x);
        for(auto j : dfn) res[j] = ORD[!use[i]].y;
    }
    return res;
}

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;
    ORD[0] = {A, 1}; ORD[1] = {B, 2}; ORD[2] = {C, 3};
    sort(ORD, ORD+3);

    UF.clear();
    for(int i=0; i<M; i++){
        int s = p[i], e = q[i];
        if(UF.merge(s, e)) Tree[s].push_back(e), Tree[e].push_back(s);
        G[s].push_back(e); G[e].push_back(s);
    }

    int cent = get_cent(0, get_sz(0));
    for(auto i : Tree[cent]){
        auto dfn = dfs_subtree(i, cent);
        if(dfn.size() >= ORD[0].x) return get_answer(dfn);
    }

    UF.clear();
    for(auto i : Tree[cent]) dfs_merge(i, cent);
    for(int i=0; i<M; i++){
        int s = UF.find(p[i]), e = UF.find(q[i]);
        if(s == cent || e == cent || s == e) continue;
        Comp[s].push_back(e); Comp[e].push_back(s);
    }

    memset(chk, 0, sizeof chk);
    for(int i=0; i<N; i++){
        if(i == cent || i != UF.find(i) || chk[i]) continue;
        auto dfn = dfs_comp(i);

        set<int> st;
        int sum = 0;
        for(auto j : dfn){
            sum += UF.W[j];
            st.insert(j);
            if(sum >= ORD[0].x) break;
        }
        if(sum < ORD[0].x) continue;

        vector<int> res;
        for(int j=0; j<N; j++) if(st.count(UF.find(j))) res.push_back(j);
        return get_answer(res);
    }
    return vector<int>(N, 0);
}

#ifdef LOCAL
int main(){
    int n, m, a, b, c;
    assert(5 == scanf("%d%d%d%d%d", &n, &m, &a, &b, &c));
    vector<int> p(m), q(m);
    for(int i=0; i<m; i++) assert(2 == scanf("%d%d", &p[i], &q[i]));
    vector<int> result = find_split(n, a, b, c, p, q);
    for (int i=0; i<(int)result.size(); i++) printf("%s%d", ((i>0)?" ":""), result[i]);
    printf("\n");
    return 0;
}
#endif

Compilation message (stderr)

split.cpp: In function 'std::vector<int> find_split(int, int, int, int, std::vector<int>, std::vector<int>)':
split.cpp:117:23: warning: comparison of integer expressions of different signedness: 'std::vector<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
  117 |         if(dfn.size() >= ORD[0].x) return get_answer(dfn);
      |                       ^

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