[백준 6603] 로또 with Python

문제 링크

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

6603번: 로또

풀이

import sys

def print_combination(n, S):
    T = [False for _ in range(len(S))]
    C = []
    C_lst = []
    def get_combination():
        if len(C) == n:
            C_lst.append(" ".join(map(str, C)))
        else:
            for i in range(len(S)-1, -1, -1):
                if T[i]:
                    i += 1
                    break
            for j in range(i, len(S)):
                if T[j]: continue
                T[j] = True
                C.append(S[j])
                get_combination()
                T[j] = False
                C.pop()

    get_combination()
    print("\n".join(C_lst))

new_line = False

while True:
    case = sys.stdin.readline()
    if int(case.split()[0]) == 0: break
    if new_line: print()
    S = list(map(int, case.split()[1:]))
    print_combination(6, S)
    new_line = True

print_combination은 get_combination을 통해 C_lst에 가능한 조합들을 모두 담아 출력한다.

 

get_combination의 동작은 다음과 같다.

base case: C의 개수가 n개(이 문제에서는 6개)가 되면, C_lst에 조합을 문자열로 추가한다.

else: 첫 번째 반복문에서 T 중에서 가장 마지막 True의 index를 찾고, i는 가장 마지막 True 바로 다음 수의 index를 가지게 된다. (모두 False면 i=0) 이 반복문은 조합이 오름차순이 되도록 하기 위함이다.

두 번째 반복문에서 T[j]가 False일 때(아직 조합에 채택되지 않은 상태) T[j]를 True로 바꾸고, 그에 해당하는 수(S[j])를 C에 추가한다. 그리고 나서 get_combination을 재귀호출하고, 그 이후에 T[j]를 False로 바꾸고 S[j]를 C에서 제거하며 다시 S[j]가 선택되지 않은 상태로 되돌리고 반복을 계속한다.

여기서 C에 숫자를 넣고 재귀호출하는 과정에서 C에 담긴 수가 6개가 되면 base case가 실행되어 조합이 C_lst에 저장된다.

 

 

이 문제는 itertools 모듈을 사용하면 더 편리하게 해결할 수 있다.

itertools를 활용하여 print_combination을 아래와 같이 바꿀 수 있다.

import itertools

def print_combination(n, S):
    comb = itertools.combinations([str(x) for x in S], n)
    print("\n".join(map(" ".join, comb)))