[프로그래머스] 해시 > 소수 만들기

소수 만들기

문제

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

• 입력
  • nums : 숫자들이 들어있는 배열
• 출력
  • 소수가 되는 경우의 개수

 

제한사항

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

 

예시

nums result

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

 

예시 설명

• 예제 #1
  • [1,2,4]를 이용해서 7을 만들 수 있습니다.
• 예제 #2
  • [1,2,4]를 이용해서 7을 만들 수 있습니다.
  • [1,4,6]을 이용해서 11을 만들 수 있습니다.
  • [2,4,7]을 이용해서 13을 만들 수 있습니다.
  • [4,6,7]을 이용해서 17을 만들 수 있습니다.

 

구현 아이디어

1. 에라토스테네스의 체를 이용하여 소수 판별하는 함수 만들기
2. 조합 함수를 이용하여 3가지을 뽑아 합을 만들 수 있는 모든 경우 탐색

 

코드

from math import sqrt
from itertools import combinations

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

def solution(nums):
    answer = 0
    rep = list(combinations(nums, r = 3))
    
    for i in rep:
        if is_prime_num(sum(i)):
            answer += 1    

    return answer

 

문법 정리

• itertools.combinations(iterable, r)
  • 입력 iterable 에서 요소의 r 길이 부분 조합 리스트를 반환
  • 조합 튜플은 iterable 입력의 순서에 따라 사전순으로 방출
  • r : 몇 개의 조합으로 출력할 지정

• 예제 1.

from itertools import combinations

a = [1, 2, 3]
print(list(combinations(a, r=2)))

# 결과
[(1, 2), (1, 3), (2, 3)]

• 예제 2.

from itertools import combinations

a = [1, 2, 3, 4, 5, 6]
print(list(combinations(a, r=5)))

# 결과
[(1, 2, 3, 4, 5), (1, 2, 3, 4, 6), (1, 2, 3, 5, 6), (1, 2, 4, 5, 6), 
 (1, 3, 4, 5, 6), (2, 3, 4, 5, 6)]