[Programmers] Lv1. 소수 만들기

문제

바로가기

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

 

풀이

• primeNumber()는 에라토스테네스의 체를 이용해 소수리스트를 만드는 함수로, 에라토스테네스의 체 연습용
• solution()은 combinations라이브러리를 사용해 nums에서 3개를 뽑고, 그 3개의 합이 소수리스트에 들어있으면 +1

 

코드

# 에라토스테네스의 체
def primeNumber():
    lst, prime = [0] * 3001, []
    lst[0] = lst[1] = 1
    for i in range(2, 3001):
        if not lst[i]:
            prime.append(i)
            for j in range(2 * i, 3001, i):
                lst[j] = 1
    return prime    

def solution(nums):
    from itertools import combinations
    answer = 0
    prime_lst = primeNumber()
    for item in combinations(nums, 3):
        if sum(item) in prime_lst:
            answer += 1
    return answer