on1ystar 머릿속을 정리, 기록하는 곳

Python itertools를 이용한 순열, 조합, 중복 순열, 중복 조합

02 Mar 2021


알고리즘 문제를 풀다 보면 순열이나 조합을 사용해야 할 때가 자주 있다. 이를 구현하기 위해서는 보통 재귀함수를 사용해야 하며, 반복문으로 구현할 시 좀 복잡한 인덱스 계산이 필요하다. 하지만 python에는 표준 라이브러리에 함수형 프로그래밍 분류로 itertools가 있는데, 여기에 순열과 조합 함수가 구현되어 있다. 속도도 빠르므로 알고 있기만 한다면 사용하는 편이 좋다.

필자는 항상 까먹고 찾아보기 때문에 공식 문서 내용과 간단한 백준 예제를 함께 정리한다.


🔷 순열(Permutation)

🔹 itertools.permutations(iterable, r=None)

iterable에서 요소의 연속된 길이 r 순열을 반환

r이 지정되지 않았거나 None이면, r의 기본값은 iterable의 길이

입력 iterable이 정렬되어 있으면, 순열 튜플이 정렬된 순서로 생성

>>> permutations('ABCD', 2)
AB AC AD BA BC BD CA CB CD DA DB DC
>>> permutations(range(3))
012 021 102 120 201 210

🔹 BOJ 15649번 - N과 M (1)

https://www.acmicpc.net/problem/15649

from itertools import permutations

N, M = map(int, input().rstrip().split())
for p in list(map(" ".join, permutations(map(str, range(1, N+1)),M))):
    print(p)


🔷 조합(combination)

🔹 itertools.combinations(iterable, r)

내용은 permutations와 같음

단, 요소는 값이 아니라 위치로 고유성을 다룬다(값이 같더라도 위치가 다르면 다른 요소로 판단).

>>> combinations('ABCD', 2)
AB AC AD BC BD CD
>>> combinations(range(4), 3)
012 013 023 123
>>> combinations('ABCA', 2)
AB AC AA BC BA CA

🔹 BOJ 15650번 - N과 M (2)

https://www.acmicpc.net/problem/15650

from itertools import combinations

N, M = map(int, input().rstrip().split())
for p in list(map(" ".join, combinations(map(str, range(1, N+1)),M))):
    print(p)


🔷 중복순열(permutation with repetition)

https://ko.wikipedia.org/wiki/중복순열

n 개의 서로 다른 원소 중에서 중복을 허용하여 r개를 뽑아서 한 줄로 나열하는 경우의 수

r 개를 선택하는 경우, 최초에 n 개를 선택할 수 있고 이후에도 계속 n 개를 선택할 수 있기 때문에 이 순열의 개수는 n^r

🔹 itertools.product(*iterables, repeat=1)

곱(product) 튜플이 정렬된 순서로 반환

원래는 데카르트 곱(Cartesian Product)을 위한 함수로 사용하지만(이름도 product) repeat 키워드 인자를 사용하여 반복 횟수를 지정하면 iterable 자신과의 곱을 할 수 있다.

>>> product('ABCD', 'xy')
Ax Ay Bx By Cx Cy Dx Dy
>>> product(range(2), repeat=3)  # product(range(2), range(2), range(2))와 동일
000 001 010 011 100 101 110 111

🔹 BOJ 15651번 - N과 M (3)

https://www.acmicpc.net/problem/15651

from itertools import product

N, M = map(int, input().rstrip().split())
for p in list(map(" ".join, product(map(str, range(1, N+1)),M))):
    print(p)

참고로 아래는 재귀로 구현한 코드

from sys import stdout

def main(permutation=[]):
    if len(permutation) == M:
        stdout.write(" ".join(permutation) + "\n")
        permutation.pop()
        return permutation
    for i in range(N):
        permutation.append(num_li[i])
        permutation = main(permutation)
    if permutation:
        permutation.pop()
        return permutation


if __name__ == "__main__":
    N, M = map(int, input().rstrip().split())
    num_li = []
    num_li = list(map(str, range(1, N+1)))
    main()


🔷 중복조합(combination with repetition)

https://ko.wikipedia.org/wiki/조합

서로 다른 n개의 원소에서 중복을 허락하여 k개를 뽑는 경우의 수

공식은 위 링크에서 참고

🔹 itertools.combinations_with_replacement(iterable, r)

combinations와 같지만, 개별 요소를 두 번 이상 반복할 수 있음

>>> combinations_with_replacement('ABC', 2)
AA AB AC BB BC CC
>>> combinations_with_replacement('ABA', 2)
AA AB AA BB BA AA

🔹 BOJ 15652번 - N과 M (4)

https://www.acmicpc.net/problem/15652

from sys import stdout
from itertools import combinations_with_replacement

N, M = map(int, input().rstrip().split())
for c in map(" ".join, combinations_with_replacement(map(str, range(1,N+1)), M)):
    stdout.write(c + "\n")

아래는 재귀로 구현한 코드

def main(combi=[], start_i=0):
    if len(combi) == M:
        stdout.write(" ".join(combi) + "\n")
        combi.pop()
        return combi
    for i in range(start_i, N):
        combi.append(num_li[i])
        combi = main(combi, i)
    if combi:
        combi.pop()
        return combi


if __name__ == "__main__":
    N, M = map(int, input().rstrip().split())
    num_li = []
    num_li = list(map(str, range(1, N+1)))
    main()


REFERENCE

https://docs.python.org/ko/3/library/itertools.html#itertool-functions

Related posts