파이썬 순열, 조합 스택으로 구현하기

보통 순열이나 조합을 구현할 때는 재귀함수를 이용한 dfs로 구현을 하게 됩니다.

하지만 저에게는 재귀함수보다는 stack을 이용한 구현이 좀 더 직관적이고 이해가 쉽기때문에 stack을 이용한 구현방법을 살펴보겠습니다.

 

순열

lst = [1,2,3,4]
n = 2

permutations = []
stack = [([x], [i]) for i, x in enumerate(lst)]  # 순열저장리스트, 인덱스리스트
print('초기상태 스택 :', stack)

while stack:
    per, idx_list = stack.pop()
    
    # n개를 모두 뽑은 경우 순열을 추가한 후 continue
    if len(per)==n:
        permutations.append(per)
        continue
        
    # n개를 뽑지않은경우
    for i in range(len(lst)):
        # 뽑은 인덱스리스트에 포함되지 않은 경우 추가
        if i not in idx_list:
            stack.append((per+[lst[i]], idx_list+[i]))

print('만들어진 순열리스트')
for p in sorted(permutations):
    print(p)

  처음 스택에 들어가 있는 값은 ([x], [i])입니다. [x]는 점차 원소들을 뽑으면서 만들어지는 순열을 저장하는 리스트입니다. [i]는 지금까지 뽑은 원소들의 인덱스들을 저장하여 중복된 인덱스의 원소를 뽑는 것을 막는 역할을 하는 리스트입니다.

  이후 스택이 있는경우 반복하면서 n개를 모두 뽑은 경우 순열리스트에 추가한 후 continue로 이후를 스킵합니다. 아직 뽑아야 할 원소가 있는 경우 원소들을 순환하며 해당 원소의 인덱스가 중복되지 않는다면 stack에 1개의 원소를 더 뽑은 튜플 (per+[lst[i]], idx_list+[i])을 추가합니다.

# 출력결과
초기상태 스택 : [([1], [0]), ([2], [1]), ([3], [2]), ([4], [3])]
만들어진 순열리스트
[1, 2]
[1, 3]
[1, 4]
[2, 1]
[2, 3]
[2, 4]
[3, 1]
[3, 2]
[3, 4]
[4, 1]
[4, 2]
[4, 3]

 

조합

lst = [1,2,3,4]
n = 2

combinations = []
stack = [([x], i) for i, x in enumerate(lst)]  # 조합저장리스트, 마지막인덱스값
print('초기상태 스택 :', stack)

while stack:
    comb, last_idx = stack.pop()
    
    # n개를 모두 뽑은 경우 조합을 추가한 후 continue
    if len(comb)==n:
        combinations.append(comb)
        continue
    
    # n개를 뽑지않은경우 마지막인덱스+1부터 추가
    for i in range(last_idx+1, len(lst)):
        stack.append((comb+[lst[i]], i))
            
print('만들어진 조합리스트')
for p in sorted(combinations):
    print(p)

  조합 구현시에는 인덱스리스트가 필요하지 않습니다. 필요한건 마지막으로 뽑은 원소의 인덱스 값이므로 초기 스택에는 조합을 저장하는 [x]와 마지막 인덱스 값을 저장하는 i만 추가합니다.

  그리고 while문에서는 n개를 뽑지 않았을 경우 마지막 인덱스+1 값에서부터 원소를 탐색합니다. 스택에는 1개를 더 뽑은 조합인 comb+[lst[i]]와 마지막 인덱스 값 i를 추가합니다.

 

# 출력결과
초기상태 스택 : [([1], 0), ([2], 1), ([3], 2), ([4], 3)]
만들어진 조합리스트
[1, 2]
[1, 3]
[1, 4]
[2, 3]
[2, 4]
[3, 4]