[Python] 소수 만들기

문제 설명

주어진 숫자 중 3개의 수를 더했을 때 소수가 되는 경우의 개수를 구하려고 합니다. 숫자들이 들어있는 배열 nums가 매개변수로 주어질 때, nums에 있는 숫자들 중 서로 다른 3개를 골라 더했을 때 소수가 되는 경우의 개수를 return 하도록 solution 함수를 완성해주세요.

 

제한사항

• nums에 들어있는 숫자의 개수는 3개 이상 50개 이하입니다.
• nums의 각 원소는 1 이상 1,000 이하의 자연수이며, 중복된 숫자가 들어있지 않습니다.

 

입출력 예

nums	result
[1, 2, 3, 4]	1
[1, 2, 7, 6, 4]	4

 

코드 작성

from itertools import combinations

def solution(nums):
    answer = 0
    temp_list = list(combinations(nums, 3))
    num_list = []
    for i in temp_list:
        num_list.append(sum(i))
    num_list = num_list
    max_num = max(num_list)
    
    num = set(range(2, max_num + 1))
    for i in range(2, int((max_num ** 0.5)) + 1):
        if i in num:
            num -= set(range(2 * i, max_num + 1, i))
    for i in num_list:
        if i in num:
            answer += 1
    return answer

- 조합 함수인 combinations를 이용하면 nums 배열에 담긴 숫자를 3개 꺼내어 합한 수를 모두 구할 수 있다.

- 2부터 가장 큰 숫자에 +1한 값까지의 모든 값을 num_list에 저장한다.

- 아래 부분은 아리토스테네스의 체를 이용한 알고리즘이다.

- 만들 수 있는 가장 큰 수의 제곱근까지의 수까지만 증가하는 for 반복문을 이용한다.

  그리고 2부터 차례대로 본인의 숫자의 배수들을 num set에서 제거해 나가면 된다.

- 모두 제거하고 난 num set은 소수만 담겨져 있는 배열이 된다.

- 이전에 만들었던 수의 조합인 num_list의 원소들이 소수배열인 num set에 포함되는지 확인한다.