[프로그래머스/python] 소수 만들기

문제

[프로그래머스] 소수 만들기 (Level 1)

프로그래머스

 

---

풀이

 

1. 주어진 리스트에 대해 3개의 숫자를 중복 없이 조합할 수 있는 모든 경우의 수를 만든다.

2. 조합의 합이 소수인지 아닌지 판별한다.

  - 소수라면, count + 1

  - 소수가 아니라면, 넘어간다.

 

위의 규칙을 구현하여 count를 반환해주면 되는 문제로, 어렵지 않게 구현을 할 수 있었다.

 

풀이 방법을 정리하고 가장 먼저 떠오른 방법은, 파이썬의 itertools.combinations 모듈을 활용하는 것이었다.

주어진 인자의 모든 조합을 튜플로 만들어서 반환해주는 굉장히 편리한 모듈이다.

반환받은 튜플들은 아래 예시처럼 list()로 변환시키면 for문 등으로 편리하게 사용할 수 있다.

 

list(combinations('ABCD', 2))
list(combinations([1, 2, 3, 4], 3))

# 결과 :
# [('A', 'B'), ('A', 'C'), ('A', 'D'), ('B', 'C'), ('B', 'D'), ('C', 'D')]
# [(1, 2, 3), (1, 2, 4), (1, 3, 4), (2, 3, 4)]

 

[CODE]

from itertools import combinations


# 소수 판별 함수
def is_prime(n):
    for i in range(2, n):
        if n % i == 0:
            return False
    return True


def solution(nums):
    answer = 0
    comb = list(combinations(nums, 3))  # combination 사용

    for temp in comb:  # temp => (1, 2, 3), (1, 2, 4), (1, 3, 4)... 
        if is_prime(sum(temp)):
            answer += 1  # 소수일 때 +1

    return answer

 

---

 

파이썬은 위처럼 모듈을 활용해서 짧고 간단하게 구현할 수 있지만,

combination 같은 모듈이 없는 환경에서는 모든 경우의 수를 탐색하는 알고리즘을 직접 구현해야 할 것이다.

 

 

그래서 겸사겸사, DFS를 이용해서 문제 풀이를 해보았다.

 

[CODE] (DFS)

arr = []  # dfs로 탐색하면서 숫자를 저장시켜놓는다.
count = 0


def dfs(_nums, idx, _len):
    global count

    if _len == 3:  # arr에 원소가 3개일 때
        if is_prime(sum(arr)):
            count += 1
        return

    for i in range(idx, len(_nums)):
        if _len < 3:
            arr.append(_nums[i])  # 현재 숫자를 arr에 넣어준다
            dfs(_nums, i+1, _len+1)
            arr.pop()  # 맨 오른쪽 숫자를 제거해주면서, 모든 경우의 수를 탐색할 수 있게 된다.


# 소수 판별 함수
def is_prime(n):
    for i in range(2, n):
        if n % i == 0:
            return False
    return True


def solution(nums):
    global count
    count = 0

    dfs(nums, 0, 0)
    return count


if __name__ == "__main__":
    print(solution([1, 2, 3, 4]))  # 출력: 1
    print(solution([1, 2, 7, 6, 4]))  # 출력: 4

 

 

동작 원리는 dfs 함수로 3자리의 모든 조합을 만들어 낸다는 것 외에는, combination 모듈을 사용했을 때와 다르지 않다.