백준 1017 <소수 쌍> Python

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

---

소수 체크와 이분매칭을 합쳐놓은 문제입니다.

우선 각 리스트의 요소에서 더하면 소수가 되는 목록을 만들어 노드간의 연결을 구해줍니다.

그 뒤 이분매칭을 해주면 됩니다.

전체에서 전체로 이분매칭을 하는 것이 왜 문제의 해답이 되는가 라고 생각할 수 있습니다.

문제에서 제시된 첫 번째 예시를 예로 들겠습니다. 

connected_with_first : 첫 번째 수가 연결된 수 / connected : 선택된 수가 연결되어있는 인덱스

    connected_with_first : 4 / connected : [3, 0, 1, 2, 5, 4]
    connected_with_first : 10 / connected : [3, 2, 1, 0, 5, 4]
    connected_with_first : 12 / connected : [3, 2, 1, -1, 5, 0]

 

전체에서 전체로 이분매칭을 돌릴 경우 이와 같이 결과가 나오는데,

1     ->    10

4     ->     1

7     ->     4

10   ->     7

11    ->    12 

12    ->    11

첫 번째 수가 10 과 연결된 경우 / 첫 번째 수가 4 와 연결된 경우

 

1     ->    10

4     ->     7

7     ->     4

10   ->     1

11    ->    12 

12    ->    11

두 경우 모두 첫 번째 수가 10과 연결된 경우로 같은 경우입니다.

 

모든 수가 매칭될 수 있을 경우, 어떤 수와 매칭된 수를 제거하는 과정을 거쳐도 남은 수들이 모두 매칭되므로 어차피 원하는 경우도 모두 매칭되게 됩니다. 그러므로 매칭되는 경우의 수를 구하는 경우 해당 방법으로 풀면 별도의 과정을 거쳐야 하지만 단순히 첫 번째 수와 연결되는 수를 찾을 때는 고르는 수 / 골라지는 수를 나누지 않아도 풀 수 있습니다.

 

다만 합이 소수라는 점에서 두 리스트를 홀수 / 짝수로 나누어 이분매칭을 할 수도 있습니다.

짝수인 소수는 2 밖에 없는데 2는 1+1 밖에 없기 때문에 중복되는 수가 없는 조건에서 존재할 수 없습니다.

즉 (홀수 + 홀수) 와 (짝수 + 짝수)는 고려의 대상이 되지 않기 때문에 위와 같이 나눌 수 있습니다. 또 이 방법이 좀 더 정석적인 풀이라고 볼 수 있겠습니다.

---

import sys

input = sys.stdin.readline

# 소수 찾기

def find_prime(num):

    # 소수인 수 체크

    prime_check = [True for _ in range(num*2 + 1)]

    # 소수 리스트

    primes = []

    # 순회

    for n in range(2, num*2 + 1):

        # 체크되지 않았다면

        if prime_check[n]:

            # 소수 추가

            primes.append(n)

            # 해당 수의 배수 모두 체크

            for i in range(n, num*2 + 1, n):

                prime_check[i] = False

    return primes

# 이분매칭

def bimatch(idx, connectable, visited, connected):

    # 이미 체크한 수이면

    if visited[idx]:

        return False

    # 체크

    visited[idx] = True

    # 연결될 수 있는 수 목록 순회

    for node in connectable[idx]:

        # 연결될 수 있는 수 이거나 다른 수와 연결할 수 있으면

        if connected[node] < 0 or bimatch(connected[node], connectable, visited, connected):

            # 연결

            connected[node] = idx

            return True

       

    return False

def solution(N, num_list):

    # 리스트에 들어있는 수 중 최대인 수의 두배보다 작은 소수 모두 찾기

    primes = find_prime(max(num_list))

    # 리스트의 각 원소에 더해서 소수가 될 수 있는 원소 인덱스 리스트

    connectable = [[] for _ in range(N)]

    # 리스트 순회

    for i in range(N-1):

        for j in range(i+1, N):

            # 두 수의 합이 소수면

            if num_list[i] + num_list[j] in primes:

                # 각 원소에 더해서 소수가 될 수 있는 원소 인덱스 추가

                connectable[i].append(j)

                connectable[j].append(i)

    # 출력할 결과

    first_connected = []

    # 첫 번째 수와 연결될 수 있는 수로 미리 임의 매칭

    for i in range(len(connectable[0])):

        connected = [-1 for _ in range(N)]

        connected[connectable[0][i]] = 0

        # 순회

        for n in range(N):

            # 방문 리스트

            visited = [False for _ in range(N)]

            # 첫 번째 수는 변경하지 않음

            visited[0] = True

            # 이분매칭

            bimatch(n, connectable, visited, connected)

        # 모든 수가 매칭되면

        if sum([1 for c in connected if c >= 0]) == N:

            # 출력할 수에 추가

            first_connected.append(num_list[connectable[0][i]])

    if first_connected:

        print(*sorted(first_connected))

    else:

        print(-1)        

       

# 입력

N = int(input().strip())

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

solution(N, num_list)