[프로그래머스 level2 소수만들기] itertools를 활용하고 이전에 풀었던 문제를 활용해보자

오늘 해결할 문제는 이전에 풀었던 문제들과 굉장히 연관이 있다는 것을 알았다.

'소수만들기'라는 말을 보고 이전에 소수 찾기 문제를 풀었던 것이 기억났다 !

 

오늘 해결해야하는 문제는 다음과 같다.

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

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

 

 

[문제 설명]

주어진 숫자 중 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

 

 

 

3개의 수를 더했을 때를 가정하였기 때문에, 과거에 순열과 조합을 처리하던 itertools가 떠올랐다. 더하는 경우, 숫자의 순서는 상관이 없기 때문에 조합을 사용해야겠다는 생각을 하였다.

따라서, itertools에서 조합은 combinations이므로 이를 사용하여 list 형태로 저장하였다.

 

소수찾기 문제에서 가장 인상깊었던 풀이는 2부터 특정 숫자까지 '집합'형태로 넣어두고, 2부터 특정숫자까지 순회하면서  해당 수가 집합에 있는 경우, 그의 배수들을 집합에서 삭제하는 형태를 통해 소수를 찾았었다. (소수는 1과 자기자신 외에 약수를 가지지 않기 때문이다 ❗️)

 

이것을 응용하여 3개의 수를 뽑아서 더하는 경우는 주어진 전체 수를 더한 것 보다 작을 것이므로, 주어진 전체 수를 더한 범위에 존재하는 소수를 찾았다. 

 

소수를 찾고 숫자 3개의 조합으로 이루어진 리스트를 더한 것이 소수라면 answer += 1을 해주는 형식으로 코드를 마무리하였다.

 

 

이러한 아이디어로 작성한 코드는 다음과 같다.

import itertools
def solution(nums):
    answer = 0
    combi_list = list(itertools.combinations(nums,3))
    sum_nums = sum(nums)+1
    sosu = set(range(2,sum_nums))
    for i in range(2,sum_nums) :
        if i in sosu :
            sosu -= set(range(2*i,sum_nums,i))
    for i in range(len(combi_list)) :
        temp = sum(combi_list[i])
        if temp in sosu :
            answer += 1
    return answer

 

 

오늘 itertools를 다시 돌아볼 수 있었고, 소수찾기에서 사용했던 알고리즘을 돌아볼 수 있었던 계기가 된 것 같다.

앞으로 활용할 일이 있다면 바로바로 접목할 수 있었으면 좋겠다 😊