[프로그래머스] 소수 만들기 - 파이썬(Python)

문제 링크

https://programmers.co.kr/learn/courses/30/lessons/12977

코딩테스트 연습 - 소수 만들기

 

문제 설명

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

 

제한사항

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

 

입출력 예

nums	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을 만들 수 있습니다.

 

문제 풀이 및 코드

문제를 접근할 때 소수인지 판별하는 코드와 세 가지 수를 더하는 코드를 혼합하여 적절히 이용하면 좋겠다고 생각하였습니다.

효율적으로 소수를 찾기 위해 math 라이브러리와 세 가지 수를 더하는 경우의 수를 찾기 위한 itertools 라이브러리를 사용하기로 결정한 후 코드를 작성하기 시작했습니다.

 

소수 판별하기

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

math 라이브러리를 통해 제곱근을 구하는 math.sqrt() 함수를 사용합니다. math.sqrt() 함수는 해당하는 숫자의 양의 제곱근을 구하는 함수입니다.

is_prime_number() 함수에 들어간 num의 제곱근을 사용하여 2부터 num의 제곱근까지 나누어 떨어지는 수가 존재하면 num은 소수가 아님을 판별하는 코드입니다.

 

세 가지 수를 더하는 경우의 수 찾기

import itertools
 
nums = [1, 2, 3, 4]
 
for i in itertools.combinations(nums, 3):
    print(i)

itertools 라이브러리의 itertools.combinations() 함수는 입력받은 배열의 경우의 수를 찾는 함수입니다. 여기서 중요한 것은 combinations는 순열과 조합에서 조합을 뜻하는 것으로 순서를 중요시하지 않는 것이 특징입니다.

[1, 2, 3, 4] 배열에서 [1, 2, 3], [1, 3, 2], [2, 1, 3], [2, 3, 1], [3, 1, 2], [3, 2, 1] 등의 결과는 모두 같은 것으로 취급한다는 것입니다.

따라서 [1, 2, 3, 4]의 itertools.combinations() 결과는 [1, 2, 3], [1, 2, 4], [1, 3, 4], [2, 3, 4] 네 가지 결괏값 만을 출력하게 됩니다.

중복된 것을 모두 포함하여 출력하고 싶을 경우 순열 함수 itertools.permutations() 함수를 이용합니다.

 

최종 코드

prime_number 배열에 세 가지 수를 더하는 경우의 수를 구하여 담은 뒤 소수 판별하는 함수를 통해 소수를 판별합니다.

소수로 판정된 개수를 결괏값에 더하여 반복한 뒤 최종 결괏값을 반환합니다.

import math
import itertools
 
 
def is_prime_number(num):
    for i in range(2, int(math.sqrt(num)) + 1):
        if num % i == 0:
            return False
    return True
 
def solution(nums):
    answer = 0
    prime_number = []
 
    for i in itertools.combinations(nums, 3):
        prime_number.append(sum(i))
 
    for i in prime_number:
        if is_prime_number(i):
            answer += 1
        else:
            continue
 
    return answer

본 코드는 제가 직접 작성한 것으로 더 나은 코드가 존재할 수 있으니 참고하시기 바랍니다.