[백준] 1017 소수 쌍 (C++)

링크: https://www.acmicpc.net/problem/1017

문제

지민이는 수의 리스트가 있을 때, 이를 짝지어 각 쌍의 합이 소수가 되게 하려고 한다. 예를 들어, {1, 4, 7, 10, 11, 12}가 있다고 하자. 지민이는 다음과 같이 짝지을 수 있다.

1 + 4 = 5, 7 + 10 = 17, 11 + 12 = 23

또는

1 + 10 = 11, 4 + 7 = 11, 11 + 12 = 23

수의 리스트가 주어졌을 때, 지민이가 모든 수를 다 짝지었을 때, 첫 번째 수와 어떤 수를 짝지었는지 오름차순으로 출력하는 프로그램을 작성하시오. 위의 예제에서 1 + 12 = 13으로 소수이다. 그러나, 남은 4개의 수를 합이 소수가 되게 짝지을 수 있는 방법이 없다. 따라서 위의 경우 정답은 4, 10이다.

입력

첫째 줄에 리스트의 크기 N이 주어진다. N은 50보다 작거나 같은 자연수이며, 짝수이다. 둘째 줄에 리스트에 들어있는 수가 주어진다. 리스트에 들어있는 수는 1,000보다 작거나 같은 자연수이며, 중복되지 않는다.

출력

첫째 줄에 정답을 출력한다. 없으면 -1을 출력한다.

---

풀이

1) 간선 연결

먼저 모든 수를 짝짓기 전에 각 수가 짝지을 수 있는 수들을 조사해야 한다.

리스트에 들어있는 수는 1000 이하이므로, 두 수의 합은 최대 1999이다.

따라서 1999까지 소수의 목록을 구한 뒤, 두 수의 합이 소수에 해당되는지를 판별해야 한다.

 

예제 1에서 각 수가 짝지을 수 있는 수들은 다음과 같다.

1 → 4, 10, 12

4 → 1, 7

7 → 4, 10, 12

10 → 1, 7

11 → 12

12 → 1, 7, 11

 

2) 이분 매칭으로 경우의 수 조사

이제 모든 수를 짝지어 첫 번째 수와 짝지어지는 수들의 집합을 구해보자.

만약 모든 경우의 수를 조사하려 한다면, 이는 시간 초과로 이어질 것이다.

50개의 수가 있고 각 수마다 2개의 짝지을 수 있는 수가 있다 하더라도 2^50만큼을 조사해야 하기 때문이다.

 

그렇기에 이분 매칭을 사용해야 하는데, 아래와 같은 방법으로 경우의 수를 조사할 수 있다.

1. 첫 번째 수(a)와 짝지을 수 있는 수(b) 하나를 고정

2. a와 b를 제외한 수들을 이분 매칭으로 모두 매칭

3. 만약 모든 수를 매칭했다면 정답에 b를 추가

4. a와 짝지을 수 있는 모든 수 b에 대하여 1~3을 반복

 

dfs를 활용한 이분 매칭의 시간 복잡도는 O(VE)이므로 여유롭게 시간 초과를 피할 수 있다.

 

코드

#include <iostream>
#include <vector>
#include <set>
using namespace std;

int nums[51];
bool isPrime[2001];

vector<int> isAble[51];
int assign[51];
bool visited[51];
set<int> answerSet;

// 1~2000까지의 소수 조사
void init(){
    for (int i = 2; i <= 2000; i++){
        isPrime[i] = true;
    }


    isPrime[0] = false;
    isPrime[1] = false;
    for (int i = 2; i <= 2000; i++){
        if (isPrime[i]){
            int j = i * 2;
            while (j <= 2000){
                isPrime[j] = false;
                j += i;
            }
        }
    }
}

// 각 수마다 짝지을 수 있는 수 조사
void findAble(int n){
    for (int i = 1; i <= n; i++){
        for (int j = i+1; j <= n; j++){
            if (isPrime[nums[i]+nums[j]]){
                isAble[i].push_back(j);
                isAble[j].push_back(i);
            }
        }
    }
}

void resetVisited(int n){
    for (int i = 1; i <= n; i++){
        visited[i] = false;
    }
}

void resetAssign(int n){
    for (int i = 1; i <= n; i++){
        assign[i] = 0;
    }
}

// 이분 매칭
bool dfs(int idx, int a, int b){
    if (assign[idx]){
        return true;
    }
    for (int i = 0; i < isAble[idx].size(); i++){
        int target = isAble[idx][i];

        if (visited[target] || target == a || target == b){
            continue;
        }

        visited[target] = true;
        int origin = assign[target];

        if (origin == 0){
            assign[idx] = target;
            assign[target] = idx;
            return true;
        }
        else {
            assign[target] = 0;
            assign[origin] = 0;
            if (dfs(origin, a, b)){
                assign[idx] = target;
                assign[target] = idx;
                return true;
            }
            else{
                assign[target] = origin;
                assign[origin] = target;
            }
        }
    }
    return false;
}

// 첫 번째 수와 짝지을 수 있는 모든 경우에 대하여 이분 매칭
void solve(int n){
    for (int i = 0; i < isAble[1].size(); i++){
        resetAssign(n);
        int target = isAble[1][i];
        assign[1] = target;
        assign[target] = 1;

        bool flag = true;
        for (int j = 2; j <= n; j++){
            resetVisited(n);
            if (!dfs(j, 1, target)){
                flag = false;
                break;
            }
        }

        if (flag){
            answerSet.insert(nums[target]);
        }
    }
}

int main(){
    ios::sync_with_stdio(false);
    cin.tie(NULL); cout.tie(NULL);

    int N;
    cin >> N;

    for (int i = 1; i <= N; i++){
        cin >> nums[i];
    }

    init();

    findAble(N);

    solve(N);

    if (answerSet.empty()){
        cout << -1;
    }
    else{
        for (auto i: answerSet){
            cout << i << " ";
        }
    }
}