[백준/1017] 소수 쌍 (Python)

Problem : https://www.acmicpc.net/problem/1017

1017번: 소수 쌍

 

Difficulty : Platinum 3

 

Status : Solved

 

Time : 00:32:57

 

---

 

문제 설명

 

 

지민이는 수의 리스트가 있을 때, 이를 짝지어 각 쌍의 합이 소수가 되게 하려고 한다. 예를 들어, {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을 출력한다.

 

입력 예시

 

6
1 4 7 10 11 12

 

출력 예시

 

4 10

 

 

---

 

풀이

 

1. 수의 합이 최대 10000이고, 최대 50^2가지 조합이 가능하므로 합의 조합을 일일히 소수 판정하던, 에라스토테네스의 체를 사용하던 소수 쌍을 구하도록 하자.

 

2. 그 다음은 가능한 소수 쌍을 구해보아야 한다. 브루트포스로 일일히 구하려면 경우의 수가 최대 O(N!)로 매우 커지게 되므로, 좀 더 적절한 방식을 사용해 볼 법하다. 이분 매칭 개념을 시도해 보는 것도 좋은 방법일 것이다.

 

3. 이분 매칭 알고리즘은 그리디하게 최대 유량을 배정할 수 있는 알고리즘이다. 즉 이 알고리즘으로 최대 유량을 구하고, 최대 유량이 N과 동일하다면 모든 정수쌍이 소수가 될 수 있는 경우인 셈이다. 그렇지 않다면 소수 쌍으로 나타내지 못하는 경우가 생김을 의미한다.

 

4. 또한 이분 매칭 알고리즘을 조금 수정할 필요가 있다. 첫째로, 우리는 첫 번째 수와 매칭되는 경우의 수를 알고 싶으므로 첫 번째 노드, 그리고 연결된 노드에 대해서는 수정이 불가능하도록 코드가 이루어져야 한다. 또한 dfs로 매칭이 되는 경우, 한 쌍을 이루는 두 노드가 연결된 상태를 유지해야 한다. 따라서 matching[node] = pair, matching[pair] = node 로 정의할 필요가 있고, 이미 매칭된 노드의 경우 탐색할 필요가 없으므로 이를 검사하는 코드도 필요하다.

 

 

풀이 코드

N = int(input())
num_list = list(map(int, input().split()))

def is_prime(n) :
  if n == 2 :
    return True
  for i in range(2, int(n ** 0.5) + 1) :
    if n % i == 0 :
      return False
  return True

prime_list = [list() for _ in range(N)]

for i in range(N-1) :
  for j in range(i+1, N) :
    if is_prime(num_list[i] + num_list[j]) :
      prime_list[i].append(j)
      prime_list[j].append(i)

def bipartite_matching(idx) :
  matching = [-1]*N
  matching[0] = idx
  matching[idx] = 0

  def dfs(node) :
    if visited[node] :
      return False
    visited[node] = True
    for pair in prime_list[node] :
      if matching[pair] == -1 :
        matching[pair] = node
        matching[node] = pair
        return True
    for pair in prime_list[node] :
      if matching[pair] not in [0, idx] and dfs(matching[pair]) :
        matching[pair] = node
        matching[node] = pair
        return True
    return False

  result = 2
  for i in range(1, N) :
    if matching[i] == -1 :
      visited = [False]*N
      if dfs(i) :
        result += 2
  return result == N 

result = list()
for i in prime_list[0] :
  if bipartite_matching(i) :
    result.append(num_list[i])

if result :
  print(*sorted(result))
else :
  print(-1)

풀이 완료!