[Python] 조합 알고리즘 구현하기

백준 2798번을 풀다 조합을 이용한 풀이를 생각했고, 이전에 모듈로 Combination을 가져다 쓴 기억이 있어 한번 직접 만들어 보는 것도 좋겠다 싶어 Combination을 구하는 알고리즘을 만들어 보았다.

만약 빠르게 조합을 구현하고 싶다면 itertools 라이브러리를 import 해서 조합, 순열 등을 구현할 수 있다. 
이는 다음 블로그에 잘 정리되어 있다
https://yganalyst.github.io/etc/memo_18_itertools/

[Python] itertools, 원소의 경우의 수(순열, 조합) 추출하기

• 조합의 정의는 다음과 같다. (출처 : https://ko.wikipedia.org/wiki/%EC%A1%B0%ED%95%A9)

 

조합(Combination)은 서로 다른 n개의 원소를 가지는 어떤 집합 (사실, 집합은 서로 다른 원소의 모임으로 정의된다.)에서 순서에 상관없이 r개의 원소를 선택하는 것이며, (즉, 선택의 순서와 상관없이 같은 원소들이 선택되었다면 같은 조합이며 다른 원소들이 선택되었다면 다른 조합이다.) 이는 n개의 원소로 이루어진 집합에서 r개의 원소로 이루어진 부분집합을 만드는 것 혹은 찾는 것과 같다.

 

• 이때 n개의 조합에서 r개의 부분집합을 고르는 조합의 경우의 수는 다음과 같다.

< nCr 공식 >

 

하지만 그 경우의 수들을 전부 보고 싶다면 어떻게 해야 할까?

단순히 중첩 반복문으로 구현하는 생각을 할 수 있겠지만, r의 값이 달라짐에 따라 r 중 반복문을 구현해야 하므로 어려워진다.

 

< 우리가 원하는 조합을 보는 방법 >

 

여기서 필자는 2진수에서 비트를 올림 하는 것처럼 숫자를 하나하나 다음 자리로 올려 주는 방법을 생각했다.

< 조합 알고리즘 설명 >

• 다음은 5C3을 구현하는 알고리즘의 일부이다.
• 1을 끝자리로 갈 때까지 다음 자리로 옮겨 주며(시프트), 만약 끝자리로 이동하면 그다음 1을 옮기며, 이때 그 뒤에 있는 1은 옮긴 1 뒤로 옮겨 준다(초기화).

 

 

이를 구현한 알고리즘은 다음과 같다.

< 조합 알고리즘 표 >

 

 

전체 코드는 다음과 같다. (깃헙 - https://github.com/siejwkaodj/Combination)

입력) 첫 줄에 정수 n, r 입력

# Combination 수동 구현
import math
n, r = map(int, input().split())

select = [1] * r + [0] * (n - r)
cnt = 0
for i in range(math.factorial(n) // (math.factorial(n - r) * math.factorial(r))):
    cnt += 1
    print(select, cnt)
    for j in range(n - 2, -1, -1):  # n-2는 끝에서 두 번째 자리부터 시작한다는 뜻.
        if select[j]:               # 현재 칸에 1이 있을 경우 다음 칸에 1이 없으면 시프트. 마지막 칸은 가지 않으니 괜찮음(이래서 n-2에서 시작)
            if select[j + 1] == 0:  # 시프트
                select[j] = 0
                select[j + 1] = 1
                if j + 2 <= n - 1 and 1 in select[j+2:]:    # 리스트 길이 검사, 길이 넘는다 해도 두 번째 조건은 건너뜀.
                    l = select[j+2:].count(1)               # 초기화 전 저장
                    select[j+2:] = [0] * len(select[j+2:])  # 뒷 자리 초기화(시프트 한 자리 바로 앞으로 이동시킴)
                    for k in range(l):                      # 뒤에 있는 1의 개수만큼 1 넣어줌
                        select[j+2+k] = 1
                break

 

 

완성된 프로그램을 이용한 예시는 다음과 같다.

< 7C3을 계산한 결과 >

 

추가사항 및 제작시 어려웠던 점)

- 팩토리얼을 구현하기 위해 math 라이브러리를 사용했다.

- 시프트를 할 때 마지막 자리가 넘어가지 않도록 n-2번째에서부터 앞 칸으로 반복한다.

- 처음에는 시프트를 하고 나서 그 뒤의 1들을 시프트 한 1 앞으로 옮겨야 한다는 걸 잊어버려 구현이 안 되던 버그가 있었다.

- j+1 자리로 시프트를 한 다음, 초기화를 위해 j+2번째 자리부터 1이 있는지 검사해야 하는데 이때 리스트의 길이를 넘지 않기 위해 미리 j+2 <= n-1을 검사한다. 파이썬은 앞 조건이 틀리면 and연산자에서는 뒷 조건을 보지 않아 길이를 넘었다 해도 오류가 뜨진 않는다.

- 뒷 자리 초기화 시 뒷자리에 있는 1의 위치를 받아 그 자리만 0으로 바꿔주기보다 그냥 뒷자리를 통째로 0으로 초기화하고 1을 시프트 자리 뒤에 붙이는 게 더 간단할 것 같아 그렇게 했다.

- r이 n보다 클 경우 다음과 같은 팩토리얼을 계산할 때 다음과 같은 오류가 난다.