[프로그래머스] 소수 만들기 - 파이썬

링크: https://school.programmers.co.kr/learn/courses/30/lessons/12977

문제 설명

주어진 숫자 중 3개의 수를 더했을 때 소수가 되는 경우의 개수를 구하려고 합니다. 숫자들이 들어있는 배열 nums가 매개변수로 주어질 때, nums에 있는 숫자들 중 서로 다른 3개를 골라 더했을 때 소수가 되는 경우의 개수를 return 하도록 solution 함수를 완성해주세요.

제한사항

• nums에 들어있는 숫자의 개수는 3개 이상 50개 이하입니다.
• nums의 각 원소는 1 이상 1,000 이하의 자연수이며, 중복된 숫자가 들어있지 않습니다.

---

입출력 예

nums	result
[1,2,3,4]	1
[1,2,7,6,4]	4

입출력 예 설명

입출력 예 #1

[1,2,4]를 이용해서 7을 만들 수 있습니다.

입출력 예 #2

[1,2,4]를 이용해서 7을 만들 수 있습니다.

[1,4,6]을 이용해서 11을 만들 수 있습니다.

[2,4,7]을 이용해서 13을 만들 수 있습니다.

[4,6,7]을 이용해서 17을 만들 수 있습니다.

문제풀이

소수 확인하는 방법:

• 어떤 수에 대해 2부터 자기 자신-1 까지의 모든 수들을 확인하며 만약 그 수로 나눠떨어지면 소수 X, 나눠 떨어지지 않으면 소수이다.

위의 방법은 하나의 수에 대해 판별하는 방법이고 위 문제는 3가지 수의 조합에 대해 찾는 것이기 때문에 조금 응용해서 for 문 세 개를 돌며 가능한 조합에 대해 소수를 확인하는 방법을 사용해야 한다.

그래서 다음과 같이 코드를 작성할 수 있다.

def solution(nums):
    answer = 0

    for i in range(len(nums)):  # 조합할 첫 번째 수
        for j in range(i+1, len(nums)):  # 조합할 두 번째 수
            for k in range(j+1, len(nums)):  # 조합할 세 번째 수
                is_prime = nums[i] + nums[j] + nums[k]  # 소수인지 확인할 수
                prime = True  
                for factor in range(2, is_prime):  # 2부터 자기자신 - 1 까지의 숫자를 돌면서
                    if is_prime % factor == 0:  # 약수로 자기자신이 나눠지면 소수가 아니다.
                        prime = False
                if prime:  # prime이 True면 정답을 1 올림
                    answer += 1
                    

    return answer

하지만 이렇게 작성한 코드는 세 개의 for 문을 돌면서 조합을 찾고 그 조합에 대해 소수인지 판별하는 for 문 까지 총 4개로 시간복잡도가 $O(n^4)$ 이다. 그래서 약수의 성질(특정 자연수의 모든 약수를 찾을 때 가운데 약수(제곱근)까지만 확인하면 된다)을 활용해서 시간복잡도를 조금이라도 줄여보려고 한다.

• ex) 16의 약수는 1, 2, 4, 8, 16
  • 제곱근 4 를 기준으로 16을 만드는 수가 대칭되기 때문에 4 까지만 확인하면 됨.

최종코드

def solution(nums):
    answer = 0

    for i in range(len(nums)):  # 조합할 첫 번째 수
        for j in range(i+1, len(nums)):  # 조합할 두 번째 수
            for k in range(j+1, len(nums)):  # 조합할 세 번째 수
                is_prime = nums[i] + nums[j] + nums[k]  # 소수인지 확인할 수
                prime = True  
                for factor in range(2, int(is_prime**0.5)+1):  # 2부터 자기자신 - 1 까지의 숫자를 돌면서
                    if is_prime % factor == 0:  # 약수로 자기자신이 나눠지면 소수가 아니다.
                        prime = False
                if prime:  # prime이 True면 정답을 1 올림
                    answer += 1
                    

    return answer

이렇게 하면 시간 복잡도를 $O(n^{3.5})$ 까지는 줄일 수 있다.

더욱 혁신적으로 줄일 수 있는 에라토스테네스의 체 라는 방법도 있지만 나중에 다뤄보도록 하겠다.