[프로그래머스] 소수 만들기 (파이썬, 최적화)

[문제]

주어진 숫자중에 3개를 더했을 때 소수가 되는 경우의 수를 세는 문제였다.

 

 

 

[기존 코드]

숫자 3개를 뽑아서 더하고 리스트에 추가한 다음, 리스트에 있는 숫자들이 소수인지 판별해주었다.

i == 3 일 때는 별도로 처리해주었다.

def solution(nums):
    # 숫자 3개 뽑아서 더하기
    t = len(nums)
    lst = []
    for i in range(t):
        for j in range(i+1,t):
            for k in range(j+1,t):
                lst.append(nums[i]+nums[j]+nums[k]) 
                
    # 소수 판별
    answer = 0
    for i in lst:
        if i == 3:
            answer += 1
        else:
            p = 2
            while p <= (i // 2):
                if i % p == 0:
                    break
                elif p == i // 2:
                    answer += 1
                    break
                else:
                    p += 1
    return answer

 

 

 

[바꾼 코드]

나는 소수를 판별할 때 $i // 2$ 까지 while문을 돌렸지만, $\sqrt{i}$까지만 돌려도 되었었다.

$\sqrt{i}$까지만 확인하면 이후에는 대칭적으로 반복되기 때문이다!

시간복잡도가 $O(N)$에서 $O(\sqrt{N})$으로 줄어든다.

import math

def solution(nums):
    # 숫자 3개 뽑아서 더하기
    t = len(nums)
    lst = []
    for i in range(t):
        for j in range(i+1, t):
            for k in range(j+1, t):
                lst.append(nums[i] + nums[j] + nums[k]) 

    # 소수 판별 (범위를 √i로 최적화)
    answer = 0
    for i in lst:
        if i == 3:
            answer += 1
        else:
            p = 2
            while p <= int(math.sqrt(i)):  # 범위를 루트i로 변경
                if i % p == 0:
                    break
                p += 1
            else:
                answer += 1
    return answer

 

 

 

[느낀점]

시간복잡도를 최대한 줄이는 방향으로 이리저리 생각하면서 코드를 짜보자.

그리고 문제를 다양하게 풀어보면서 경험을 쌓는 것도 중요한 것 같다.