(프로그래머스 스택/큐 문제 풀이) 소수 찾기

문제 설명
한자리 숫자가 적힌 종이 조각이 흩어져있습니다. 흩어진 종이 조각을 붙여 소수를 몇 개 만들 수 있는지 알아내려 합니다.
각 종이 조각에 적힌 숫자가 적힌 문자열 numbers가 주어졌을 때, 종이 조각으로 만들 수 있는 소수가 몇 개인지 return 하도록 solution 함수를 완성해주세요.

제한사항
numbers는 길이 1 이상 7 이하인 문자열입니다.
numbers는 0~9까지 숫자만으로 이루어져 있습니다.
"013"은 0, 1, 3 숫자가 적힌 종이 조각이 흩어져있다는 의미입니다.

입출력 예
numbers           return
"17"                    3
"011"                  2

입출력 예 설명
예제 #1
[1, 7]으로는 소수 [7, 17, 71]를 만들 수 있습니다.

예제 #2
[0, 1, 1]으로는 소수 [11, 101]를 만들 수 있습니다.
- 11과 011은 같은 숫자로 취급합니다.

from itertools import permutations
import math

# 소수 판별 함수
def check(n): 
    k = math.sqrt(n) # 제곱근까지만 봐도 소수 판별 가능
    if n < 2: # 1은 소수가 아님
        return False

    for i in range(2, int(k) + 1): # 16을 예로 들면, (2, 8), (4, 4), (8, 2), 즉 (8, 2)까지 가지 않고 4까지만 보면 됨
        if n % i == 0:
            return False
    return True

def solution(numbers):
    answer = []
    for n in range(1, len(numbers) + 1): # 각 원소 개수로 만들 수 있는 모든 경우의 수
        per = list(map(''.join, permutations(list(numbers), n))) # n개 원소 조합으로 수열 만들기
        for i in list(set(per)): # 반복 숫자 set으로 제거
            if check(int(i)):
                answer.append(int(i))

    answer = len(set(answer)) # 각 원소 개수로 만든 숫자에서 발생할 수 있는 반복 제거
    return answer

소수를 판별하는 문제로, 주어진 numbers에서 만들 수 있는 모든 가능한 경우의 수를 구해 소수 여부를 판단하면 된다. 파이썬에서 모든 가능한 경우의 수를 구하는 것은 itertools 라이브러리를 활용하면 쉽다.

permutation 함수의 기본 형식
permutations(list, n))) <- n개 원소로 조합

permutations 테스트 결과를 보면 훨씬 이해하기 쉽다. 주의할 점은 반복 가능한 리스트 내부 원소는 문자열 형태여야 한다.

from itertools import permutations

test = ['1', '2', '3']
print(list(map(''.join, permutations(test, 3)))) # 3개 원소 모두 사용

permutations 자체를 출력하면 기본적으로 객체를 반환한다. 그래서 for 문으로 원소를 살펴보면 아래 형태와 같다. 위에서 map 함수로 각 결과값에 ''.join을 해주는 이유이다.

print(permutations(test, 3))

for i in permutations(test, 3):
  print(i)

그런데 보통 코딩 테스트에서 itertools 라이브러리를 사용하게 해주겠지? 라는 생각이 든다. ㅎㅎ