백준 1154 - 팀 편성

백준 1154 - 팀 편성

풀이

한 학생을 기준으로 해당 학생과 아는 사이가 아닌 학생들은 서로 같은 그룹이어야 한다.
이를 union-find를 사용하여 그룹을 묶는다.

각 그룹들의 학생들이 서로 모두 아는 사이인지 확인한다.
이때, 해당 그룹은 여러 개가 나올 수 있는데 1번 학생이 포함된 그룹에 포함시킬 수 있는 만큼 포함시키고 나머지는 다른 그룹에 합친다. 다른 그룹들에 대해서 모두 아는 사이인지 확인하면 된다.

주의할 점은 어느 한 팀에는 인원이 없을 수 있다.

코드

#include <bits/stdc++.h>

#define endl "\n"
#define all(v) (v).begin(), (v).end()
#define For_IMPL(condition, i, a, b, increment, decrement) \
    for (int i = (a); condition; (a < b ? increment : decrement))
#define For(i, a, b) For_IMPL((a < b ? i < b : i > b), i, a, b, ++i, --i)
#define For_(i, a, b) For_IMPL((a < b ? i <= b : i >= b), i, a, b, ++i, --i)
#define ft first
#define sd second

using namespace std;
using ll = long long;
using lll = __int128_t;
using ulll = __uint128_t;
using ull = unsigned long long;
using ld = long double;
using pii = pair<int, int>;
using pll = pair<ll, ll>;
using ti3 = tuple<int, int, int>;
using tl3 = tuple<ll, ll, ll>;

template<class T> bool ckmin(T& a, const T& b) { return b < a ? a = b, 1 : 0; }
template<class T> bool ckmax(T& a, const T& b) { return a < b ? a = b, 1 : 0; }

const int INF = 987654321;
const int INF0 = numeric_limits<int>::max();
const ll LNF = 987654321987654321;
const ll LNF0 = numeric_limits<ll>::max();

const int mxn = 1e3+1;

int n;
bool adj[mxn][mxn];
int p[mxn];

int f(int i) {
    if(i == p[i]) return i;
    return p[i] = f(p[i]);
}
bool uni(int a, int b) {
    a = f(a), b = f(b);
    if(a == b) return false;
    p[b] = a;
    return true;
}

bool check(vector<int> &v) {
    for(auto a : v) {
        for(auto b : v) {
            if(!adj[a][b]) return false;
        }
    }
    return true;
}
bool check2(vector<int> &a, vector<int> &b) {
    for(auto aa : a) {
        for(auto bb : b) {
            if(!adj[aa][bb]) return false;
        }
    }
    return true;
}

void solve() {
    cin >> n;
    memset(adj, 0, sizeof adj);
    For_(i,0,n) adj[i][i] = true;
    while(true) {
        int a, b; cin >> a >> b;
        if(a==-1 and b==-1) break;
        adj[a][b] = adj[b][a] = true;
    }
    iota(p,p+mxn,0);
    queue<int> q;
    For_(i,1,n) {
        For_(j,1,n) {
            if(adj[i][j]) continue;
            q.emplace(j);
        }
        int u;
        if(!q.empty()) {u = q.front(); q.pop();}
        while(!q.empty()) {
            int v = q.front(); q.pop();
            uni(u,v);
        }
    }

    map<int,vector<int>> ans;
    For_(i,1,n) ans[f(i)].emplace_back(i);

    for(auto x : ans) {
        if(!check(x.second)) {
            cout << "-1\n";
            return;
        }
    }

    vector<int> team[2];
    for(auto x : ans[f(1)]) {
        team[0].emplace_back(x);
    }
    for(auto x : ans) {
        if(x.ft == f(1)) continue;
        if(check2(team[0], x.sd)) {
            for(auto y : x.sd) team[0].emplace_back(y);
        }
        else {
            for(auto y : x.sd) team[1].emplace_back(y);
        }
    }

    if(!check(team[1])) {
        cout << "-1\n";
        return;
    }

    cout << "1\n";
    sort(all(team[0]));
    for(auto x : team[0]) cout << x << ' ';
    cout << "-1\n";
    sort(all(team[1]));
    for(auto x : team[1]) cout << x << ' ';
    cout << "-1\n";

}

int main(void) {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);

    int tc = 1;
//    cin >> tc;
    while(tc--) {
        solve();
//        cout << solve() << endl;
    }


    return 0;
}

후기

문제의 테스트 케이스가 부족한 듯 하다. 내가 생각한 반례가 통과하는 경우도 종종 있어서 내 풀이 또한 반례가 있을 수 있다.
또한 그룹을 합치는 과정에서 최악의 경우 $O(N^3)$이 나올 수도 있다고 생각하는데 그런 테스트 케이스가 없는 건지 아니면 그렇게 최악까지는 안 가는 건지 확신은 못 하겠다.