[프로그래머스][Lv1][Python] 소수 만들기

[문제]

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

코딩테스트 연습 - 소수 만들기

 

[코드]

def check(a, b, c):
    total = a + b + c
    for i in range(2, total // 2 + 1):
        if total % i == 0: return False
    return True

def solution(nums):
    answer = 0
    for i in range(0, len(nums) - 2):
        for j in range(i + 1, len(nums) - 1):
            for k in range(j + 1, len(nums)):
                if check(nums[i], nums[j], nums[k]): answer += 1
    return answer

- solution 함수의 for문 3개를 통해 nums 리스트에서 나올 수 있는 모든 숫자 조합을 구한다.

 

- check 함수를 통해 주어진 숫자 3개를 더했을 때의 숫자가 소수인지 확인한다.

   - 소수를 구할 때 for문의 end를 해당 숫자를

     2로 나눈 몫까지로 지정해 필요 없는 범위 계산

     을 줄였다.

 

 

 

코드에서 예상했듯이 첫 번째 풀이의 시간은 꽤 오래걸린다..

for문을 세개나 중첩했다는 것부터 문제가 있다. 따라서 개선이 필요한 코드라고 할 수 있다.

 

[코드 개선]

from itertools import combinations 

def check(a, b, c):
    total = a + b + c
    for i in range(2, total // 2 + 1):
        if total % i == 0: return False
    return True

def solution(nums):
    answer = 0
    for i in list(combinations(nums, 3)): 
        if check(i[0], i[1], i[2]): answer += 1
    return answer

모든 경우를 for문 3개가 아닌 다른 방법으로 구하는게 제일 중요하다. 이는 itertools의 combinations 메소드로 해결할 수 있다.

 

combinations

리스트에 있는 값들의 모든 조합을 구할 수 있는 라이브러리인 itertools에는 permutations, combinations, product가 있는데 각 차이점은

• 하나의 리스트에서 모든 조합 구하기
  
  • permutations
  • combinations

items = ['1', '2', '3', '4', '5']

from itertools import permutations list(permutations(items, 2))
# [('1', '2'), ('1', '3'), ('1', '4'), ('1', '5'), ('2', '1'), ('2', '3'), ('2', '4'), ('2', '5'), ('3', '1'), ('3', '2'), ('3', '4'), ('3', '5'), ('4', '1'), ('4', '2'), ('4', '3'), ('4', '5'), ('5', '1'), ('5', '2'), ('5', '3'), ('5', '4')]

from itertools import combinations list(combinations(items, 2))
# [('1', '2'), ('1', '3'), ('1', '4'), ('1', '5'), ('2', '3'), ('2', '4'), ('2', '5'), ('3', '4'), ('3', '5'), ('4', '5')]

• 두개 이상의 리스트의 모든 조합 구하기
  
  • product

from itertools import product

items = [['a', 'b', 'c,'], ['1', '2', '3', '4'], ['!', '@', '#']]
list(product(*items))
# [('a', '1', '!'), ('a', '1', '@'), ('a', '1', '#'), ('a', '2', '!'), ('a', '2', '@'), ('a', '2', '#'), ('a', '3', '!'), ('a', '3', '@'), ('a', '3', '#'), ('a', '4', '!'), ('a', '4', '@'), ('a', '4', '#'), ('b', '1', '!'), ('b', '1', '@'), ('b', '1', '#'), ('b', '2', '!'), ('b', '2', '@'), ('b', '2', '#'), ('b', '3', '!'), ('b', '3', '@'), ('b', '3', '#'), ('b', '4', '!'), ('b', '4', '@'), ('b', '4', '#'), ('c,', '1', '!'), ('c,', '1', '@'), ('c,', '1', '#'), ('c,', '2', '!'), ('c,', '2', '@'), ('c,', '2', '#'), ('c,', '3', '!'), ('c,', '3', '@'), ('c,', '3', '#'), ('c,', '4', '!'), ('c,', '4', '@'), ('c,', '4', '#')]

 

테스트 코드를 돌려보면 개선 전 코드보다 itertools를 import해 구한 코드로 시간이 꽤 줄어들은 것을 볼 수 있다.