[프로그래머스 2단계]소수_만들기

---

문제 설명

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

제한사항

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

입력

1. nums : [1,2,3,4]
2. nums : [1,2,7,6,4]

출력

1. result : 1
2. result : 4

문제 풀이

• 우선 소수는 2를 제외하면 무조건 홀수이기 때문에 세 숫자를 더해서 홀수가 되는 경우만 고려했습니다.
• (숫자 셋을 더하면 무조건 2이상이므로 2를 예외처리 하지 않았습니다.)
• 소수 판별은 내장 함수를 추가해 판별합니다.

def is_prime(x):
    for i in range(2, int(math.sqrt(x))+1):
        if x % i == 0:
       		return False
    return True

• 주어진 nums를 홀수 odds와 짝수 even으로 구분합니다. 방법의 개수를 저장할 변수 count도 지정해줍니다.

odds = [n for n in nums if n%2 != 0]
even = [n for n in nums if n%2 == 0]
count = 0

• (1) 홀수 + 홀수 + 홀수

itertools의 combinations를 이용하여 홀수 3개로 만들 수 있는 모든 조합을 구해 그 중 숫자 셋을 더해 소수가 되는 경우 count를 올립니다.

if len(odds) >= 3:
    for three_odd in itertools.combinations(odds, 3):
        if is_prime(sum(three_odd)):
            count += 1

• (2) 홀수 + 짝수 + 짝수
  even에서 2개의 짝수를 뽑아 그 둘의 조합의 합과 odds의 홀수 1개와 순차적으로 더해가며 소수를 판별합니다.

if len(odds) >= 1 and len(even) >= 2:
    for odd in odds:
        for two_even in itertools.combinations(even,2):
            if is_prime(sum(two_even)+odd):
                count += 1

최종 제출 코드

import itertools
import math
def solution(nums):
    
    def is_prime(x):
        for i in range(2, int(math.sqrt(x))+1):
            if x % i == 0:
                return False
        return True
    
    odds = [n for n in nums if n%2 != 0]
    even = [n for n in nums if n%2 == 0]
    
    #홀 홀 홀
    count = 0
    if len(odds) >= 3:
        for three_odd in itertools.combinations(odds, 3):
            if is_prime(sum(three_odd)):
                count += 1
    #홀 짝 짝
    if len(odds) >= 1 and len(even) >= 2:
        for odd in odds:
            for two_even in itertools.combinations(even,2):
                if is_prime(sum(two_even)+odd):
                    count += 1
    
    return count