[프로그래머스] 소수 만들기 / Python / 소수

소수 만들기

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

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

나의 풀이

• 소수 판별 리스트를 먼저 만듦
  - 에라토스테네스의 체를 구현할 때는 구하려는 범위의 제곱근까지만 검사하면 됨!
• nums에 있는 숫자들 중 3개를 골라 더하는 경우의 수를 combinations로 구함

import itertools

def solution(nums):
    sosu = [True] * 3001 # 1부터 3000까지 소수인지 아닌지 판별하기 위한 리스트
    sosu[0], sosu[1] = False, False # 0과 1은 소수가 아님
    
    for i in range(2, int(3000 * (1/2)) + 1): # 2부터 판별, 3000의 제곱근까지만 검사하면 됨
        if sosu[i]: # True이면
            j = 2
            while i * j <= 3000: # 그 수의 배수는 모두 False로 처리
                sosu[i * j] = False
                j += 1

    # num에 들어있는 숫자 중 3개를 골라 더하는 모든 경우의 수(조합)를 구하고
    # 3개 수의 합이 sosu 리스트에서 True 값을 가지는지 확인
    sum_lst = [sum(i) for i in list(itertools.combinations(nums, 3)) if sosu[sum(i)]]
    
    return len(sum_lst)