[algorithm][python] 프로그래머스 소수찾기

문제

소수 찾기

문제 풀이

from itertools import combinations

def solution1(nums): # 방법1 : combinations 사용
    cnt = 0
    for cb in combinations(nums, 3): # 원소 3개로 구성된 리스트의 모든 조합
        s = sum(cb)
        if isPrime(s):
            cnt += 1
    return cnt
    
def isPrime(num): # 소수 판별
    for divisor in range(2, num // 2 + 1):
        if num % divisor == 0:
            return False
    return True

print(solution1([1,2,3,4])) # 테스트코드
    

1. combinations(nums, 3) : 원소 3개로 구성된 리스트의 모든 조합을 구한다.
2. sum(cb) : 모든 원소의 합을 구한다.
3. isPrime(s) : 모든 원소의 합 s가 소수인지 검사한다.

새롭게 알게된 점

• 파이썬 기본 라이브러리중 itertools 모듈에는 combinations 함수가 있다.
• itertools.combinations(iterable, r) : iterable중에서 r개를 선택할 수 있는 조합을 이터레이터로 리턴한다.

from itertools import combinations
>>> list(combinations(items, 2))
# [('1', '2'), ('1', '3'), ('1', '4'), ('1', '5'), ('2', '3'), ('2', '4'), ('2', '5'), ('3', '4'), ('3', '5'), ('4', '5')]

combinations을 몰랐을 때는 아래와 같이 단순 반복문을 사용했었는데, 코드를 작성하면서 코드가 마음에 들지 않아 다른 풀이를 찾아봤다. 이때 combinations의 존재를 알게된 것이다. 다른 문제에서도 유용하게 잘 쓸 수 있을 것 같다.

def solution2(nums): # 방법2 : 반복문 사용
    sum = 0
    cnt = 0
    nums_len = len(nums)
    
    for i in range(nums_len - 2):
        for j in range(i+1, nums_len - 1):
            for k in range(j+1, nums_len):
                sum = nums[i] + nums[j] + nums[k]
                if isPrime(sum): # 방법1에서 사용한 소수판별 함수와 동일
                    cnt += 1
    return cnt