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

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

문제

• 정수가 담긴 리스트가 입력으로 주어졌을 때, 이 리스트의 원소 3가지를 더했을 때 이 값이 소수가 되는 경우의 수를 구하는 문제
• [1,2,3,4] 리스트의 경우 (1, 2, 4) 의 조합의 합이 7로 소수이다.

코드

from itertools import combinations


def solution(nums):
    cnt = 0

    for num in combinations(nums, 3):
        tmp = sum(num)
        for i in range(2, tmp):
            if tmp % i == 0:
                break
        else:
            cnt += 1

    return cnt

설명

• 삼총사 문제와 접근 방법은 동일하다. 리스트의 조합을 먼저 구한 후 각 조합의 합을 구한다.
• 각 조합의 합을 반복문을 돌려서 소수인지 판별한다.
• 위 코드에서 소수인지 판별하는 반복문의 범위를 제곱근의 절반으로 줄이면 속도면에서 개선될 수 있다.