[Leetcode] 리트코드 - 조합(combinations) 파이썬(python) 풀이

문제

전체의 수 n을 입력받아 k개의 조합을 리턴하라.

입력 : n = 4, k = 2
출력

[
    [1,2],
    [1,3],
    [1,4],
    [2,3],
    [2,4],
    [3,4]
]

코드

dfs

from typing import List


class Solution:
    def combine(self, n: int, k: int) -> List[List[int]]:
        results = []

        def dfs(elements, start: int, k:int):
            if k == 0:
                results.append(elements[:])

            for i in range(start, n+1):
                elements.append(i)
                dfs(elements, i+1, k-1)
                elements.pop()

        dfs([], 1, k)
        return results

itertools 모듈

import itertools

class Solution:
    def combine(self, n: int, k: int):
        return list(itertools.combinations(range(1, n+1), k))

풀이

순열의 경우 자기 자신을 제외하고 모든 요소를 next_elements로 처리했으나, 이와 달리 조합의 경우 자기 자신뿐만 아니라 앞의 모든 요소를 배제하고 next_elements를 구성한다. 따라서 여기서는 그냥 elements라는 이름으로 처리한다.

1부터 순서대로 for 문으로 반복하되, 재귀 호출할 때 넘겨주는 값은 자기 자신 이전의 모든 값을 고정하여 넘긴다. 따라서 남아 있는 값끼리 나머지 조합을 수행하게 되며, k=0이 되면 결과에 삽입한다.

참고

print(f'range({start},{n+1}), elements: {elements}, dfs({i+1}, {k-1})')

-------------------------------------------------
range(1,5), elements: [1], dfs(2, 1)
range(2,5), elements: [1, 2], dfs(3, 0)
range(3,5), elements: [1, 2, 3], dfs(4, -1)
range(4,5), elements: [1, 2, 3, 4], dfs(5, -2)
range(3,5), elements: [1, 2, 4], dfs(5, -1)
range(2,5), elements: [1, 3], dfs(4, 0)
range(4,5), elements: [1, 3, 4], dfs(5, -1)
range(2,5), elements: [1, 4], dfs(5, 0)
range(1,5), elements: [2], dfs(3, 1)
range(3,5), elements: [2, 3], dfs(4, 0)
range(4,5), elements: [2, 3, 4], dfs(5, -1)
range(3,5), elements: [2, 4], dfs(5, 0)
range(1,5), elements: [3], dfs(4, 1)
range(4,5), elements: [3, 4], dfs(5, 0)
range(1,5), elements: [4], dfs(5, 1)

dfs 호출시마다 값을 찍어보며 이해했다.

출처
파이썬 알고리즘 인터뷰 (글 : 박상길 그림 : 정진호) [책만]