콘텐츠로 건너뛰기
Home » 순열 조합 제너레이터 구현하기

순열 조합 제너레이터 구현하기

파이썬 기본 모듈인 itertools는 리스트와 같은 반복 가능한 객체에 대한 여러 연산을 위한 함수를 제공하는데, 특정한 원소들로 만들 수 있는 순열이나 조합을 구하는 함수들도 이 모듈에 포함되어 있다. 특정한 원소를 가지고 만들 수 있는 모든 순열은 itertools.permutations() 함수로, 모든 조합은 itertools.combinations() 함수로 만들 수 있다. 그 외에 중복을 포함하는 조합등을 생성하는 함수가 제공된다.

아주아주 예전에는 순열/조합을 생성하는 함수를 구현하는 것이 간단한 작업이 아니었는데, 제너레이터에서 다른 제너레이터로 결과 생성을 위임하는 yield from 문법이 생긴 이후에는 이를 아주 간단히 구현할 수가 있게 되었다. 그래서 오늘은 제너레이터를 이용해서 주어진 리스트나 튜플에서 순열 및 조합을 생성하는 제너레이터를 만드는 방법을 소개하고자 한다.

순열을 생성하는 제너레이터

n 개의 원소를 가진 리스트나 튜플에서 k 개의 원소를 가진 부분 집합의 순열을 생성하는 제너레이터를 먼저 만들어보겠다. 사실 이 아이디어가 조합이나 다른 제너레이터를 만드는데 가장 기본이 되는 로직이 될 것이다.

만약 (a1, a2, a3, a3, ..., an) 으로 구성된 n개의 원소중에서 k를 골라 만들 수 있는 순열을 생성해야 한다면 이렇게 생각해볼 수 있을 것이다.

  1. 순열의 첫번째 원소가 될 값을 고른다. 첫번째 원소는 주어진 집합에서 순서대로 하나를 고르면 된다. 만약 첫번째 원소를 a1으로 정했다면 (a2, a3, … an)으로 구성된 n – 1 개의 남은 원소들로 길이가 k – 1 인 순열을 만들고, 각 순열의 맨 앞에 a1을 붙여주면 될 것이다.
  2. 1.의 아이디어를 약간 뒤집으면 이미 고른 x개의 원소들의 목록이 정해져 있는 상태에서, 남은 원소들로 k – x 의 길이를 가지는 순열을 만들고 이어 붙이면 된다.

따라서 재귀 함수의 형태로 재귀 제너레이터를 만들고, 미리 정해진 원소의 목록과 남은 원소의 목록을 구분하여 던져주면 된다. 이게 말로 쓰려니까 괜히 복잡하고 어렵게 느껴질 수 있는데 코드는 정말 간단하다.

def permutations(ns, k=None):
  def helper(head, tail, k):
    # head : 이미 선택한 원소들
    # tail : 아직 남은 원소드
    # k    : 남은 순열의 길이

    # k가 0이면 더 추가할 원소가 없음. 지금까지 선택한 목록을 내보냄
    if k == 0:
      yield head

    # 남은 원소 중에서 순차적으로 하나씩 선택하여, 선택된 목록으로 옮기고
    # 재귀 호출
    for i, x in enumerate(tail):
      yield from helper((*head, x), (*tail[:i], *tail[i+1:]), k-1)

  # k == None 이면 ns 전체 길이에 대한 순열을 구하는 것으로 간주
  k = k if k is not None else len(ns)
  # 선택한 목록으로는 비어있는 튜플을 전달
  yield from helper(tuple(), ns, k)

주석이 많아서 그렇지 실제로는 7줄짜리 함수 코드가 된다. 다음과 같이 [1,2,3,4,5,6] 의 모든 순열을 생성해서 출력해보자.

for i, x in permutations([1,2,3,4,5,6]):
  print(i + 1, x)

중복 순열 생성 제너레이터

중복 순열은 원래의 집합에서 같은 원소를 1개 이상 사용할 수 있는 순열이다. 따라서 재귀 호출을 할 때 방금 선택한 원소를 제외하는 부분을 빼버리면 된다.

def permutation_with_dups(ns, k=None):
  def helper(head, tail, k):
    if k == 0:
      yield head
    for i, x in enumerate(tail):
      yield from helper((*head, x), tail, k-1)
  yield from helper(tuple(), ns, k if k is not None else len(ns))

조합 생성 제너레이터

조합과 순열의 차이점은 순열은 선택한 요소들 사이의 순서가 중요하지만, 조합은 그 순서가 중요하지 않다는 것이다. 알파벳 26자 중 에서 3글자를 뽑아 순열과 조합을 만들 때, “ABC”, “BAC” 는 순열에서는 서로 다른 것이지만, 조합에서는 같은 것으로 본다. 이 부분은 의외로 간단하게 해결되는데, 앞에서부터 순서대로 원소를 선택하기 때문에 한 번 선택한 원소가 포함되는 조합 이후로는 같은 원소가 조합에 포함되지 않으면 된다.

def combinatons(ns, k=None):
  def helper(head, tail, k):
    if k == 0:
      yield head
    for i, x in enumerate(tail):
      # tail 부분에 현재 선택한 이후 부분만 전달한다.
      yield from helper((*head, x), tail[i+1:], k-1)
  yield from helper(tuple(), ns, k if k is not None else len(ns))
      

중복 조합을 생성하는 제너레이터

중복조합은 어떨까? 이전에 선택한 적이 있는 원소는 다시 선택할 수 없다. (선택하면 순열이 될 것이다.) 대신에 방금 선택한 원소는 이 다음에 다시 선택할 수 있도록 해주면 된다. 그려면 신기하게도 중복 조합을 생성하는 제너레이터가 된다.

def combinations_with_dups(ns, k=None):
  def helper(head, tail, k):
    if k == 0:
      yield head
    for i, x in enumerate(tail):
      yield from helper((*head, x), tail[i:], k-1)
  yield from helper(tuple(), ns, k if k is not None else len(ns))
      

보너스 – 중복된 원소를 포함하는 순열 만들기

순열 조합과 관련된 문제 중에는 대상 집합의 원소 중에 중복된 것이 있는 케이스가 있다. 예를 들어 “aaabbbb”로 만들 수 있는 순열은 몇 개일까? 7개 원소를 가지고 있으므로 5040개의 순열을 만들 수 있지만, 3개의 a는 모두 같고, 4개의 b는 모두 같으므로 (3! * 4!)을 나눈 35개만큼의 서로 다른 문자열을 만들 수 있다.

하지만 우리가 만든 제너레이터 함수의 아이디어를 사용하여 접근하면 의외로 간단히 해결된다. 지금까지의 아이디어의 핵심은 이미 고른 원소들과, 남아있는 고를 수 있는 원소들을 구분하여 재귀 호출의 구조에서 이 정보를 유지해 나가는 것이었다. 이 문제의 답은 아래의 글에서 알아보도록 하자.