(python)프로그래머스-소수 만들기

문제 설명
주어진 숫자 중 3개의 수를 더했을 때 소수가 되는 경우의 개수를 구하려고 합니다. 숫자들이 들어있는 배열 nums가 매개변수로 주어질 때, nums에 있는 숫자들 중 서로 다른 3개를 골라 더했을 때 소수가 되는 경우의 개수를 return 하도록 solution 함수를 완성해주세요.
제한 사항
nums에 들어있는 숫자의 개수는 3개 이상 50개 이하입니다.
nums의 각 원소는 1 이상 1,000 이하의 자연수이며, 중복된 숫자가 들어있지 않습니다.

소수란?

소수(Prime Number): 1보다 크고 1과 자기 자신을 제외한 다른 수로는 나누어지지 않는 수를 뜻한다. 소수는 등장 방식이 매우 불규칙해 종잡을 수가 없다는 특징이 있다.

에라토스테네스의 체

소수를 골라내는 가장 오래되고 유일한 방법으로 '에라토스테네스의 체'가 있다.
1) 먼저 2부터 100까지 적어둔다.
2) 2를 제외하고 2의 배수를 제거한다.
3) 3을 제외하고 3의 배수가 되는 수를 제거한다.
4) 5를 제외하고 5의 배수가 되는 수를 제거한다.
5) 7을 제외하고 7의 배수가 되는 수를 제거한다.
6) 11을 제외하고 11의 배수가 되는 수를 제거한다.
7) 이것을 반복한다. 마지막에 남은 것들이 소수이다.

이런식으로 2부터 n까지의 자연수를 나열 후 남은 수 중에서 아직 처리하지 않은 가장 작은 수 i를 찾는다. 이후 남은 수 중에서 i의 배수를 모두 제거한다.(i는 제거하지 않음.)
요 과정을 반복하는것이 에라토스테네스의 체 알고리즘이다.

import math

n = 100 # 2부터 100까지의 모든 수에 대하여 소수 판별
arr = [True for i in range(n + 1)] # 처음엔 모든 수가 소수(True)인 것으로 초기화

# 에라토스테네스의 체 알고리즘 
for i in range(2, int(math.sqrt(n)) + 1): # 2부터 n의 제곱근까지의 모든 수를 확인하며
    if arr[i] == True: # i가 소수인 경우 (남은 수인 경우)
        # i를 제외한 i의 모든 배수를 지우기
        j = 2 
        while i * j <= n:
            arr[i * j] = False
            j += 1

# 모든 소수 출력
for i in range(2, n + 1):
    if arr[i]:
        print(i, end=' ')

문제 풀이 방법
에라토스테네스의 체 알고리즘을 사용하면 시간이 많이 단축되지만 이 문제에서는 nums의 최대 길이가 50이기 때문에 그냥 for문을 돌려서 풀어주었다. combinations을 사용해 3개를 선택하는 조합을 구해주었다.
선택된 3개의 값을 모두 더한 값을 저장하는 sum_result배열을 만들어주었고 sum_result에 저장된 값들 중 소수가 있을때 answer값을 +1해주었다.

from itertools import combinations
def solution(nums):
    answer = 0
    arr_3=list(combinations(nums,3))
    sum_result=[]
    result=0
    for i in range(len(arr_3)):
        for j in range(3):
            result+=arr_3[i][j]
        sum_result.append(result)
        result=0
    
    for i in sum_result:
        count=0
        for j in range(2,i):
            if i%j==0:
                break
            count+=1
            if count==i-2:
                answer+=1
    return answer