소수 찾기(programmers 문제) python 풀이

문제 링크: https://programmers.co.kr/learn/courses/30/lessons/42839

입력으로 문자열 numbers이 들어오면 numbers 안에 있는 원소(1이상 7이하)들을 조합하여 새로운 수들을 만든다. 조합하여 만들 수 있는 수 중에서 소수의 갯수를 찾아 리턴하는 것이 목적이다.

내 풀이

1. 첫 번째

from itertools import permutations

def solution(numbers):
    answer = []
    for i in range(1, len(numbers)+1):
        map_num = list(map(''.join, permutations(list(numbers), i)))
        for j in map_num:
            if is_prime_number(int(j)):
                answer.append(int(j))
        
    answer = list(set(answer))
    
    return len(answer)

# 소수인지 아닌지 확인
def is_prime_number(number):
    if number < 2:
        return False
    elif number == 2: # 2일 경우는 소수가 맞음
        return True
    elif number % 2 == 0:
        return False
    for i in range(3, number, 2):
        if number % i == 0:
            return False
            
    return True

먼저 numbers 문자열 안에 있는 숫자들로 새로운 숫자들을 만들어야 한다.
이를 위해 itertools에 있는 premutations를 import해와서 map_num안에 i자리수의 배열을 넣어주었다. (입력이 "123"이고 i가 1이면 map_num의 값은 [1, 2, 3])

다음으로 map_num안에 있는 인덱스값들 중 소수의 갯수를 카운트해야 하기 때문에, is_prime_number이라는 함수를 따로 만들어 소수인지 아닌지 확인하고 소수라면 answer 배열에 append를 해준다.

최종적으로 배열안의 값들이 겹치면 안되므로 set과 list를 이용해 배열 내부 중복 제거를 한 뒤 배열안의 갯수를 리턴해주었다.

결과

하지만, 시간 복잡도에서 문제가 발생했다.

2. 두 번째

from itertools import permutations

def solution(numbers):
    answer = []
    for i in range(1, len(numbers)+1):
        map_num = list(map(''.join, permutations(list(numbers), i)))
        for j in map_num:
            if is_prime_number(int(j)):
                answer.append(int(j))
        
    answer = list(set(answer))
    
    return len(answer)

# 소수인지 아닌지 확인
def is_prime_number(number):
    if number < 2:
        return False
    elif number == 2:
        return True
    elif number % 2 == 0:
        return False
    # number의 약수를 루트 number까지만 확인하면 됨
    root_number = int(number ** 0.5)
    for i in range(3, root_number+1, 2):
        # 소수가 아님
        if number % i == 0:
            return False
    return True

어떠한 숫자 n의 약수 중 자기 자신을 제외한 가장 큰 약수는 무엇일까?
바로 루트 n보다 같거나 낮다는 점이다.

따라서 시간 복잡도 문제를 해결하기 위해 is_prime_number 함수 내부에 있는 for문에 number가 아닌 root_number 값을 넣어주었다.

결과

모든 테스트를 통과하였다.