[백준 1017번] 파이썬 - 소수 쌍

백준 1017_소수 쌍

시간 제한 2초, 메모리 제한 128MB

# 조건

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

# 접근 방법

• 우선 첫 번째 수의 가장 큰 경우는 N이 2인 경우 999이며, 따라서 합이 1999가 최대로 나올 수 있는 수이다.
• 에라토스테네스의 체를 이용하여 2000까지의 소수를 구해준 후 리스트에 기록해준다.

def prime_numbers():  
    for i in range(2, int(math.sqrt(2000))+1):  
        # i가 소수인 경우  
        # i를 제외한 i의 모든 배수 지우기        
        if primeNumbers[i] == True:  
            j = 2  
            while i*j <= 2000:  
                primeNumbers[i*j] = False  
                j += 1

• 이후 각 쌍을 짝지을 건데 소수가 되기 위한 성질 하나를 활용해준다.
  • 소수는 짝수와 홀수의 합으로만 이루어질 수 있다.
  • 따라서, 입력받은 숫자들을 짝수와 홀수로 나누어서 저장해준다.
• 홀수와 짝수의 수가 같지 않다면 바로 -1 출력
• 시작하는 숫자가 홀수이면 짝수 인덱스 번호를 저장해주고 아니라면 홀수 인덱스 번호를 저장해준다.
• 이분 매칭을 dfs를 활용하여 돌려줄건데, selected 리스트와 visited 리스트를 사용해준다.
  • selected 리스트는 선택되는 수가 어떤 수와 매칭되었는지를 나타내어주고
  • visited는 선택하는 수가 매칭되어 있는지 여부를 나타내어준다.
• 첫 번째 수가 매칭될 수 있는 모든 수를 탐색하며 다음 수를 bipartite_matching 함수로 돌려준다.
  • 이 때, 매칭되는 수가 선택되지 않았거나
  • 매칭한 수가 다른 수를 선택할 수 있다면 selected[idx]를 변경해준다.

단순 dfs 시간 초과

• 매칭될 수 있는 모든 숫자를 탐색하였다.

import sys, math  
sys.stdin = open('input.txt')  
input = sys.stdin.readline  

# 소수 만들기  
# 에라토스테네스의 체 알고리즘 사용  
def prime_numbers():  
    for i in range(2, int(math.sqrt(2000))+1):  
        # i가 소수인 경우  
        # i를 제외한 i의 모든 배수 지우기        
        if primeNumbers[i] == True:  
            j = 2  
            while i*j < 2000:  
                primeNumbers[i*j] = False  
                j += 1  

def bipartite_matching(idx):  
    for i in can_match[idx]:  
        if visited[i] == False:  
            visited[i] = True  
            if idx == len(odds) - 1 and not False in visited:  
                    return True  
            if idx < len(odds) - 1:  
                val = bipartite_matching(idx + 1)  
                if val:  
                    return True  
            visited[i] = False  

    return False  

N = int(input())  
nums = [*map(int, input().split())]  

# 처음엔 모든 수가 소수인 것으로 가정  
primeNumbers = [True for i in range(2000)]  
primeNumbers[0], primeNumbers[1] = False, False  
prime_numbers()  

# 입력받은 수를 짝수와 홀수로 나눠주기  
odds = []  
evens = []  

for i in nums:  
    if i % 2:  
        odds.append(i)  
    else:  
        evens.append(i)  

# 합이 소수가 되는 인덱스를 저장해주기  
flag = True  
# 첫 수가 홀수인지 짝수인지에 따라 저장  
if nums[0] % 2:  
    flag = False  
if flag:  
    can_match = [[] for _ in range(len(evens))]  
else:  
    can_match = [[] for _ in range(len(odds))]  

if not flag:  
    for i in range(len(odds)):  
        for j in range(len(evens)):  
            if primeNumbers[odds[i] + evens[j]]:  
                can_match[i].append(j)  
else:  
    for i in range(len(evens)):  
        for j in range(len(odds)):  
            if primeNumbers[odds[j] + evens[i]]:  
                can_match[i].append(j)  
result = []  
if len(odds) != len(evens):  
    print(-1)  
else:  
    for i in can_match[0]:  
        visited = [False] * len(odds)  
        if visited[i] == False:  
            visited[i] = True  
            now_val = bipartite_matching(1)  
            if now_val:  
                if flag:  
                    result.append(odds[i])  
                else:  
                    result.append(evens[i])  
        visited[i] = False  

    if result:  
        result.sort()  
        for k in result:  
            print(k, end=' ')  
    else:  
        print(-1)

# pass 코드

• selected 리스트를 활용하여 현재 매칭하려는 수를 선택한 수가, 다른 선택지가 있다면 변경해주었다.

import sys, math  
sys.stdin = open('input.txt')  
input = sys.stdin.readline  

# 소수 만들기  
# 에라토스테네스의 체 알고리즘 사용  
def prime_numbers():  
    for i in range(2, int(math.sqrt(2000))+1):  
        # i가 소수인 경우  
        # i를 제외한 i의 모든 배수 지우기        
        if primeNumbers[i] == True:  
            j = 2  
            while i*j < 2000:  
                primeNumbers[i*j] = False  
                j += 1  

def bipartite_matching(idx):  
    if visited[idx]:  
        return 0  
    visited[idx] = True  
    for k in can_match[idx]:  
        if selected[k] == -1 or bipartite_matching(selected[k]):  
            selected[k] = idx  
            visited[k] = True  
            return 1  
    return 0  


N = int(input())  
nums = [*map(int, input().split())]  

# 처음엔 모든 수가 소수인 것으로 가정  
primeNumbers = [True for i in range(2000)]  
primeNumbers[0], primeNumbers[1] = False, False  
prime_numbers()  

# 입력받은 수를 짝수와 홀수로 나눠주기  
odds = []  
evens = []  

for i in nums:  
    if i % 2:  
        odds.append(i)  
    else:  
        evens.append(i)  

# 길이 다르면 바로 탈출  
if len(odds) != len(evens):  
    print(-1)  
    exit()  
# 첫 수가 홀수인지 짝수인지에 따라 저장  
first = odds if nums[0] % 2 else evens  
second = evens if first == odds else odds  

can_match = [[] for _ in range(len(first))]  
for i,j in enumerate(first):  
    for k, u in enumerate(second):  
        if primeNumbers[j+u]:  
            can_match[i].append(k)  

result = []  
for i in can_match[0]:  
    selected = [-1] * len(second)  
    selected[i] = 0  
    cnt = 1  
    for j in range(1, len(first)):  
        visited = [False] * len(first)  
        visited[0] = True  
        cnt += bipartite_matching(j)  
    if cnt == len(second):  
        result.append(second[i])  

if result:  
    result.sort()  
    for k in result:  
        print(k, end=' ')  
else:  
    print(-1)