코딩테스트 조합 구하기 : 작은 수부터 뽑으면 쉽다

 코딩 테스트에서 많이 나오는 것 중에 하나는 조합을 구현하는 것입니다. 물론, 이것이 단독으로 나오는 경우는 매우 드뭅니다. 보통, 구현까지 요구하는 경우가 많습니다. 여기서 더 어려워 지면, 대놓고 조합으로 푸세요. 라고 하지 않고, 조건을 숨겨놓는 경우도 있는데요. 이 문제가 그에 속합니다.

 

 이번 시간에는 조합을 어떻게 구현하는지 정도만 언급하도록 하겠습니다.

 

 

---

 n개 중에 m개를 뽑는다고 해 봅시다. 뽑은 것만 중요한 게 조합입니다. 순서가 중요하진 않아요. 그러면, 하나를 뽑으면서, 뽑을 수 있는 집합을 줄여나간다고 봐도 될까요? 예를 들어보겠습니다.

 

 

 카드가 5개 있습니다. 이 중 3개를 뽑아야 합니다. 뽑은 것이 중요하니, 전략을 하나 세워봅시다. 일단, 먼저 뽑은 카드에 적혀있는 수가 나중에 뽑은 카드의 수보다 작아야 합니다. 예를 들어서, 2를 뽑은 경우에는 그 다음에 3을 뽑을 수 있지만, 반대로 2를 뽑고 1을 뽑을 수 없습니다.

 

 만약에 순서가 중요했다면, 이야기가 달라졌을 겁니다. 2를 뽑고 1을 뽑는 것과, 1을 뽑고 2를 뽑는 게 다른 경우로 취급되기 때문입니다. 하지만, 뽑은 것이 중요하지 순서는 중요하지 않으니, 전략을 수가 커지는 순으로 뽑는다로 잡아 놓은 겁니다.

 

 

 처음에 뽑을 수 있는 집합은 1, 2, 3, 4, 5가 있습니다. 이 이야기는 1을 뽑을 수도 있고, 2를 뽑을 수도 있고, 3, 4, 5를 뽑을 수도 있다는 소리입니다. 2를 뽑았다고 해 보겠습니다.

 

 

 그러면, 그 다음에는 3, 4, 5 중 하나를 뽑을 수 있습니다. 1을 뽑을 수 없는데요. 이전에 뽑았던 것이 2이므로, 2보다 작은 1을 뽑지 않습니다. 이미 전략을 수가 커지는 순으로 뽑는다고 잡아놓았는데, 2 뽑고 1을 뽑으면 전략에 위배가 되어 버립니다. 그러면, 3, 4, 5 중에 하나를 뽑을 수 있을 텐데요. 이 중에 4를 뽑았다고 해 봅시다.

 

 

 그러면, 그 다음에 뽑을 수 있는 것은 5밖에 없습니다.

 

 

 최종적으로 2, 4, 5가 뽑히게 됩니다. 여기서 핵심은 카드를 하나 뽑을 때 마다, 그 다음에 뽑을 수 있는 것의 집합이 어떻게 줄어드냐 입니다. 이것은 재귀로 구현하면 쉽습니다.

 

---

 조합을 구하는 코드를 보면서 이야기 해 보도록 하겠습니다.

 

 Combi(7, 3)을 호출합니다. 이는 7개 중에 3개를 뽑는 상황을 의미합니다.

 

 

 이 함수 안으로 들어가 봅시다. 중요한 것은 13번째 줄인데요. 1부터 n까지 for loop를 돈다는 것을 알 수 있습니다. 이는 뽑을 수 있는 번호가 [1, n]이기 때문입니다. 16번째 줄에, dfs(i+1, n, m-1)이 있는데요. 각각의 의미를 설명하면, i를 뽑은 경우에는 다음에 뽑을 수 있는 카드는 i+1 이상임을 의미합니다. 다음에 n은 뽑을 수 있는 최대 수가 n임을 의미합니다. m-1은 depth인데요.

 

 뽑을 수 있는 횟수를 의미합니다. 처음에 m개를 뽑을 수 있는데요. 하나를 뽑았다면, m-1개를 더 뽑을 수 있습니다. 계속 뽑는 횟수를 소모하다가, 0이 되면 결과만 출력하고 빠져 나오면 되는 셈입니다.

 

 

 그러면, dfs 내부를 구현하는 것도 크게 어렵지 않습니다. depth가 0이면 조합을 출력합니다. 이것은 22번째 줄 if문 블록에서 수행합니다. 그렇지 않으면, x부터 n까지 도는데요. 이는 선택하는 카드가 x인 경우, x+1인 경우, ... , n인 경우를 의미합니다.

 

 

 수행 결과는 위와 같습니다. 뽑을 수 있는 수를 어떻게 좁혀나갔는지가 핵심입니다.

 

---

 그런데, 이 마저도 귀찮다면 어떻게 할까요? 이런 메서드가 제공되는 라이브러리를 이용하면 되는데요. 파이썬에는 itertools가 있습니다. 이것을 간단하게 써 보겠습니다.

 

 

 itertools의 combinations는 길이가 r인 combination을 리턴합니다. 1번째 인자로 iterable을 받는데요. list 같은 것을 넘겨도 됩니다.

 

 

 어떻게 쓸까요? 요렇게 쓰면 됩니다. 1번째 인자에 range를 넣었는데요. 이것은 1이상 8미만을 의미합니다. 즉, 1번부터 7번까지를 의미합니다. r은 이 중에서 3개를 뽑는다는 의미입니다.

 

 

 수행 결과는 위와 같습니다.