[PS/Python] Python itertools으로 순열/조합 구하기

1. 문제 정의

프로그래머스 lv2 84512 모음사전
https://school.programmers.co.kr/learn/courses/30/lessons/84512

주어진 5개의 알파벳을 기반으로 순열 리스트를 구한 후, 완전 탐색을 이용해 주어진 단어의 위치를 찾는 문제입니다.

문제 자체는 직관적이고 쉬우나, 순열 리스트를 구하는 방법을 고민하다 python의 itertools 라는 라이브러리를 새롭게 알게 되었습니다.

2. itertools

https://docs.python.org/ko/3.8/library/itertools.html

itertools : 효율적인 루핑을 위한 이터레이터를 만드는 함수

• 무한 이터레이터:
  count(), cycle(),repeat()
• 가장 짧은 입력 시퀀스에서 종료되는 이터레이터:
  accumulate(), chain(), chain.from_iterable(), compress(), dropwhile(), filterfalse(), groupby(), islice(), startmap(), takewhile(), tee(), zip_longest()
• 조합형 이터레이터:
  product(), permutations(), combinations(), combinations_with_replacement()

다양한 이터레이터 관련 함수들을 제공하고 공식 문서에 설명이 정말 자세하게 나와있습니다. 이 중 조합형 이터레이터를 중심으로 더 알아보겠습니다.

2-1. product

itertools.product(*iterables, repeat=1)

• 인자로 여러 개의 iterables를 넣을 수 있으며, 입력 iterables의 데카르트 곱을 반환합니다.
• product(A, B)는 ((x,y) for x in A for y in B)와 동일한 결과를 반환합니다.
• iterable 자신과의 데카르트 곱을 구하려면 repeat를 이용해 곱셈 횟수를 지정할 수도 있습니다.

import itertools

letters = ['A', 'B', 'C']
nums = [1, 2]

print(list(itertools.product(letters, repeat = 2)))
# product(letters, letters, repeat = 1)와 동일한 결과
# 결과 : [('A', 'A'), ('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'B'), ('B', 'C'), ...

print(list(itertools.product(letters, repeat = 3)))
# product(letters, letters, letters, repeat = 1)와 동일한 결과
# 결과 : [('A', 'A', 'A'), ('A', 'A', 'B'), ('A', 'A', 'C'), ('A', 'B', 'A'), ('A', 'B', 'B'), ('A', 'B', 'C'), ('A', 'C', 'A'), ('A', 'C', 'B'), ('A', 'C', 'C'), ('B', 'A', 'A'), ('B', 'A', 'B'), ('B', 'A', 'C'), ...

print(list(itertools.product(nums, letters, repeat = 1)))
# 결과 : [(1, 'A'), (1, 'B'), (1, 'C'), (2, 'A'), (2, 'B'), (2, 'C')]

2-2. permutations

itertools.permutations(iterable, r=None)

• iterable에서 r개의 요소를 가지는 순열을 반환합니다.
• r이 None이면 최대 길이의 순열을 반환합니다.

import itertools

letters = ['A', 'B', 'C']

print(list(itertools.permutations(letters, 2)))
# 결과 : [('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'C'), ('C', 'A'), ('C', 'B')]

print(list(itertools.permutations(letters)))
# 결과 : [('A', 'B', 'C'), ('A', 'C', 'B'), ('B', 'A', 'C'), ('B', 'C', 'A'), ('C', 'A', 'B'), ('C', 'B', 'A')]

print(list(itertools.permutations(range(1, 4))))
# 결과 : [(1, 2, 3), (1, 3, 2), (2, 1, 3), (2, 3, 1), (3, 1, 2), (3, 2, 1)]

2-3. combinations

itertools.combinations(iterable, r)

• 입력 iterable에서 요소의 개수가 r인 조합을 반환합니다.

import itertools

letters = ['A', 'B', 'C']

print(list(itertools.combinations(letters, r = 2)))
# 결과 : [('A', 'B'), ('A', 'C'), ('B', 'C')]

print(list(itertools.combinations(letters, r = 3)))
# 결과 : [('A', 'B', 'C')]

print(list(itertools.combinations(range(1, 4), r = 3)))
# 결과 : [(1, 2, 3)]

2-4. combinations_with_replacement

itertools.combinations_with_replacement(iterable, r)

• 입력 iterable에서 요소의 개수가 r인 조합을 반환하며, 중복을 허용합니다.

import itertools

letters = ['A', 'B', 'C']

print(list(itertools.combinations_with_replacement(letters, r = 2)))
# 결과 : [('A', 'A'), ('A', 'B'), ('A', 'C'), ('B', 'B'), ('B', 'C'), ('C', 'C')]

print(list(itertools.combinations_with_replacement(letters, r = 3)))
# 결과 : [('A', 'A', 'A'), ('A', 'A', 'B'), ('A', 'A', 'C'), ('A', 'B', 'B'), ('A', 'B', 'C'), ('A', 'C', 'C'), ('B', 'B', 'B'), ('B', 'B', 'C'), ('B', 'C', 'C'), ('C', 'C', 'C')]

print(list(itertools.combinations_with_replacement(range(1, 4), r = 3)))
# 결과 : [(1, 1, 1), (1, 1, 2), (1, 1, 3), (1, 2, 2), (1, 2, 3), (1, 3, 3), (2, 2, 2), (2, 2, 3), (2, 3, 3), (3, 3, 3)]

3. 최종 코드

위에서 살펴본 조합형 이터레이터 중 itertools.product()를 이용해 문제를 해결했습니다.
순열, 조합 리스트를 구해야 하는 상황 때 itertools를 유용하게 사용해보시기 바랍니다!

from itertools import product

def solution(word):
    letters = ['A', 'E', 'I', 'O', 'U']
    alpha_list = []
    
    for i in range(1, 6):
        alpha_list += list(product(letters, repeat = i))
        
    alpha_list.sort()
    for i in range(len(alpha_list)):
        alpha_list[i] = "".join(alpha_list[i])
    
    for i in range(len(alpha_list)):
        if alpha_list[i] == word:
            return i + 1
    
    return 0