Codeforces Round 268 (Div. 1) - Two Sets 링크
(2023.07.19 기준 Difficulty *2000)

문제

n개의 모두 다른 정수가 있다. 이 정수들을 집합 A 아니면 집합 B에 넣고자 한다. 넣을 때 조건이 있는데, 정수 x가 집합 A에 들어가면 a - x도 무조건 집합 A에 있어야 한다. 또, 정수 x가 집합 B에 들어가면 b - x도 무조건 집합 B에 있어야 한다. 모든 정수를 조건에 맞게 집합 A나 B에 넣을 수 있는지 판별 및 정수별로 어느 집합에 들어가야 하는지 출력

알고리즘

2-SAT

문제

정수 별로 각 집합에 들어갈 수 있는지 검사하자.

정수 x가 집합 A에 들어가려면 a-x가 있는지부터 알아야 한다. 그러므로 주어지는 정수 x의 수열 p를 처음에 정렬해놓고 이분 탐색으로 a-x가 있는지 찾으면 된다.

만약, a-x가 있다면, x와 a-x는 무조건 같은 집합에 들어가야 한다. x + a-x = a 이므로 결국, x와 a-x는 집합 A에서 서로를 필요로 하는데, 만약에 하나가 집합 A에 들어가지 않는다면 나머지 하나는 자동으로 집합 A에 들어가지 못한다. 그러므로 x와 a-x의 관계는 xnor이다. (x xnor y) 절은 (x or ~y) and (~x or y)와 같다.

반대로 a-x가 없다면, x는 집합 A에 들어가지 못하므로 무조건 집합 B에 들어가야 한다. 그러므로 (x or x) 절을 추가하자.

이를 집합 B에도 똑같이 적용하면 된다. 나는 집합 A를 x, 집합 B를 ~x로 잡았다.

코드

• C++

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

const int MAXN = 100000;

int n, a, b;
vector<int> p;
map<int, int> origin;

int Not(int a){
    if (a <= n) return a + n;
    return a - n;
}

// strongly connected component
struct SCC{
    vector<int> graph[MAXN * 2 + 1];
    stack<int> stk;
    int ids[MAXN * 2 + 1], idx;
    int scc_ids[MAXN * 2 + 1], scc_idx;

    int _dfs(int i){
        int min_idx = ids[i] = idx++;
        stk.push(i);

        // 사이클이 도는 정점을 모두 같은 번호로
        // 사이클 중 가장 번호가 작은 값으로 scc를 이룸
        for (auto j: graph[i]){
            if (!~ids[j]) min_idx = min(min_idx, _dfs(j));
            else if (!~scc_ids[j]) min_idx = min(min_idx, ids[j]);
        }

        // 지금 번호가 사이클을 이루는 번호(가장 작은 번호)이면
        // 그 번호까지 stk에서 꺼내서 scc를 이루게 함.
        if (ids[i] == min_idx){
            // vector<int> scc;
            while (true){
                int j = stk.top(); stk.pop();
                scc_ids[j] = scc_idx;
                // scc.push_back(j);
                if (i == j) break;
            }
            // scc_lst.push_back(scc);
            scc_idx++;
        }

        return min_idx;
    }

    void tarjan(){ // get SCC
        fill(ids + 1, ids + n * 2 + 1, -1); idx = 0;
        fill(scc_ids + 1, scc_ids + n * 2 + 1, -1); scc_idx = 0;
        for (int i = 1, nn = n * 2; i <= nn; i++) if (!~ids[i]) _dfs(i);
    }

    void check(){
        int result[n];
        for (int i = 1; i <= n; i++){
            // 역과 같은 SCC에 속하면 모순
            if (scc_ids[i] == scc_ids[Not(i)]){
                cout << "NO";
                return;
            }
            // scc 번호는 위상 정렬의 역순. 위상 정렬 순서 상 빠른 것이 False
            // 즉, scc 번호가 높은 것이 False
            result[origin[p[i - 1]]] = scc_ids[i] > scc_ids[Not(i)];
        }

        cout << "YES" << '\n';
        for (int i = 0; i < n; i++) cout << result[i] << ' ';
    }
};

// two_sat
struct TS{
    SCC scc;

    void _add(int a, int b){
        scc.graph[a].push_back(b);
    }

    void add_and(int a){ // ... and a
        int _a = Not(a);
        _add(_a, a);
    }

    void add_or(int a, int b){ // ... and (a or b)
        int _a = Not(a), _b = Not(b);
        _add(_a, b); _add(_b, a);
    }

    void add_xor(int a, int b){ // ... and (a xor b) -> ... and (a or b) and (~a or ~b)
        int _a = Not(a), _b = Not(b);
        add_or(a, b); add_or(_a, _b);
    }

    void add_xnor(int a, int b){ // ... and (a xnor b) -> ... and (a or ~b) and (~a or b)
        int _a = Not(a), _b = Not(b);
        add_or(a, _b); add_or(_a, b);
    }

    void check(){ // SCC를 구한 후 모순인지 확인 및 정답 출력
        scc.tarjan();
        scc.check();
    }
}ts;

int main(){
    ios_base::sync_with_stdio(0);
    cin.tie(0);

    cin >> n >> a >> b;

    p.resize(n);
    for (int i = 0; i < n; i++) cin >> p[i], origin[p[i]] = i; // 원래의 인덱스 저장
    sort(p.begin(), p.end()); // 이분 탐색을 위하여 정렬

    bool check[n];

    // set A (x)
    fill(check, check + n, false);
    for (int i = 0; i < n; i++){
        if (check[i]) continue;
        int j = lower_bound(p.begin(), p.end(), a - p[i]) - p.begin();
        if (j < n && a - p[i] == p[j]){ // i와 j는 set A에 같이 들어가거나 같이 들어가지 않아야 한다.
            check[i] = check[j] = true;
            ts.add_xnor(i + 1, j + 1);
        }
        else ts.add_and(Not(i + 1)); // i는 set A에 들어갈 수 없으므로 무조건 set B에 들어가야 한다.
    }

    // set B (~x)
    fill(check, check + n, false);
    for (int i = 0; i < n; i++){
        if (check[i]) continue;
        int j = lower_bound(p.begin(), p.end(), b - p[i]) - p.begin();
        if (j < n && b - p[i] == p[j]){ // i와 j는 set B에 같이 들어가거나 같이 들어가지 않아야 한다.
            check[i] = check[j] = true;
            ts.add_xnor(Not(i + 1), Not(j + 1));
        }
        else ts.add_and(i + 1); // i는 set B에 들어갈 수 없으므로 무조건 set A에 들어가야 한다.
    }

    ts.check();
}

• Python (RE)

import sys; input = sys.stdin.readline
sys.setrecursionlimit(200000)
from bisect import bisect_left

def Not(a):
    if a <= n:
        return a + n
    return a - n

# strongly connected component
class SCC:
    def __init__(self):
        self.graph = [[] for _ in range(n * 2 + 1)]

    def _dfs(self, i):
        min_idx = self.ids[i] = self.idx
        self.idx += 1
        self.stk.append(i)

        # 사이클이 도는 정점을 모두 같은 번호로
        # 사이클 중 가장 번호가 작은 값으로 scc를 이룸
        for j in self.graph[i]:
            if not ~self.ids[j]:
                min_idx = min(min_idx, self._dfs(j))
            elif not ~self.scc_ids[j]:
                min_idx = min(min_idx, self.ids[j])

        # 지금 번호가 사이클을 이루는 번호(가장 작은 번호)이면
        # 그 번호까지 stk에서 꺼내서 scc를 이루게 함.
        if self.ids[i] == min_idx:
            # scc = []
            while True:
                j = self.stk.pop()
                self.scc_ids[j] = self.scc_idx
                # scc.append(j)
                if i == j:
                    break
            # scc_lst.append(scc)
            self.scc_idx += 1

        return min_idx

    def tarjan(self): # get SCC
        self.stk = []
        self.ids = [-1] * (n * 2 + 1); self.idx = 0
        self.scc_ids = [-1] * (n * 2 + 1); self.scc_idx = 0
        for i in range(1, n * 2 + 1):
            if not ~self.ids[i]:
                self._dfs(i)

    def check(self):
        result = [0] * n
        for i in range(1, n + 1):
            # 역과 같은 SCC에 속하면 모순
            if self.scc_ids[i] == self.scc_ids[Not(i)]:
                print('NO')
                return
            # scc 번호는 위상 정렬의 역순. 위상 정렬 순서 상 빠른 것이 False
            # 즉, scc 번호가 높은 것이 False
            result[origin[p[i - 1]]] = int(self.scc_ids[i] > self.scc_ids[Not(i)])

        print('YES')
        print(*result)

class TS: # two_sat
    def __init__(self):
        self.scc = SCC()

    def _add(self, a, b):
        self.scc.graph[a].append(b)

    def add_and(self, a): # ... and a
        _a = Not(a)
        self._add(_a, a)

    def add_or(self, a, b): # ... and (a or b)
        _a = Not(a); _b = Not(b)
        self._add(_a, b)
        self._add(_b, a)

    def add_xor(self, a, b): # ... and (a xor b) -> ... and (a or b) and (~a or ~b)
        _a = Not(a); _b = Not(b)
        self.add_or(a, b)
        self.add_or(_a, _b)

    def add_xnor(self, a, b): # ... and (a xnor b) -> ... and (a or ~b) and (~a or b)
        _a = Not(a); _b = Not(b)
        self.add_or(a, _b)
        self.add_or(_a, b)

    def check(self): # SCC를 구한 후 모순인지 확인 및 정답 출력
        self.scc.tarjan()
        self.scc.check()

n, a, b = map(int, input().split())
p = list(map(int, input().split()))

origin = {p[i]: i for i in range(n)} # 원래의 인덱스 저장
p.sort() # 이분 탐색을 위하여 정렬

ts = TS()

# set A (x)
check = [False] * n
for i in range(n):
    if check[i]:
        continue
    j = bisect_left(p, a - p[i])
    if j < n and a - p[i] == p[j]: # i와 j는 set A에 같이 들어가거나 같이 들어가지 않아야 한다.
        check[i] = check[j] = True
        ts.add_xnor(i + 1, j + 1)
    else: # i는 set A에 들어갈 수 없으므로 무조건 set B에 들어가야 한다.
        ts.add_and(Not(i + 1))

# set B (~x)
check = [False] * n
for i in range(n):
    if check[i]:
        continue
    j = bisect_left(p, b - p[i])
    if j < n and b - p[i] == p[j]: # i와 j는 set B에 같이 들어가거나 같이 들어가지 않아야 한다.
        check[i] = check[j] = True
        ts.add_xnor(Not(i + 1), Not(j + 1))
    else: # i는 set B에 들어갈 수 없으므로 무조건 set A에 들어가야 한다.
        ts.add_and(i + 1)

ts.check()

여담

코드포스에서의 파이썬은 재귀 제한 오류가 떠도, 재귀 제한을 늘린다고 해서 완전히 해결이 되지 않는다. 무슨 stack size가 문제라던데 thread 모듈을 이용하여 stack size를 늘려준 다음, 메인 함수를 threading 하면 된다던데.. 스택 사이즈를 늘리면 TLE, 줄이면 RE가 계속 뜬다. 아무리 조정해도. (PyPy3는 그냥 메모리 초과 뜬다..)

뭐.. 아무튼 새로운 걸 알게 되었다.