완전탐색 - Itertools

프로그래머스를 풀다가 완전 탐색 문제를 만났다.

 

위 문제는 결국 모든 조합을 찾아 소수가 되는 경우를 찾아야 한다. 전형적인 완전 탐색 문제이다. 

 

하지만 문제가 하나 더 있다. 여기서 어떻게 모든 조합을 찾아낼 수 있을까? for문을 사용해 모든 조합을 만들면 너무 코드가 복잡해진다. 이럴땐 Python의 itertools.combinations을 사용하여 간편하게 3개의 숫자를 고를 수 있다.

 

import itertools
import math

def is_prime(n):
    if n < 2:
        return False
    for i in range(2, int(math.sqrt(n)) + 1):
        if n % i == 0:
            return False
    return True

def solution(nums):
    count = 0
    # 3개의 숫자를 선택하는 모든 조합 구하기
    for comb in itertools.combinations(nums, 3):
        if is_prime(sum(comb)):  # 선택한 숫자들의 합이 소수인지 확인
            count += 1
    return count

 

이와 같이 itertools.combinations(nums, 3)을 하면 nums 배열에서 중복을 허용하지 않는 3개의 원소를 무작위로 골라 가능한 모든 조합을 만든다.