백준15656 - N과 M (7) | Python itertools 모듈

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

 

해당 문제는 등장하는 모든 수열을 중복없이 나열하는 문제로, 예전 자바스크립트로 풀던 방식처럼 백트래킹 로직을 활용하여 풀었더랬다.

import sys
data_temp = sys.stdin if sys.platform == 'linux' else open('입력.txt', 'r')
input_data = data_temp.read().splitlines()

def solution (data) :
  [_,m], temp = map(lambda x : list(map(int, x.split())),data)
  arr = sorted(temp)
  result, storage = [],[]

  def backTracking () :
    if len(storage) == m :
      result.append(' '.join(map(str, storage)))
      return

    for i in range(0, n) :
      storage.append(arr[i])
      backTracking()
      storage.pop()
  backTracking()

  print('\n'.join(result))

solution(input_data)

 

하지만 풀이 후에 다른 사람들이 풀이를 보니 itertools 모듈을 사용하여 풀이하는 것을 보았고 관심이 생겨 이에 대해 다뤄본다.

 

먼저 itertools 모듈은 파이썬 내장 모듈로 데이터를 다루는데 유용한 도구들을 제공한다. (아래 링크 참고)

해당 포스팅에서는 반복 가능한 객체(조합, 순열, 중복조합)에 대해 사용하는 

combinations | permuations | product 메서드에 대해서만 다룬다. 

 

combinations(iterable, r)

주어진 iterable 객체에서 길이가 r인 조합으로 데이터 반환 (순서 고려 X)

import itertools

iterable = ['A', 'B', 'C']
combines = itertools.combinations(iterable, 2)
print(list(combines))
'''
[
('A', 'B'), 
('A', 'C'), 
('B', 'C')
]
'''
#

 

permutations(iterable, r)

주어진 iterable 객체에서 길이가 r인 모든 순열로 데이터 반환 (순서 고려 O, 가능한 모든 순서쌍)

import itertools

iterable = ['A', 'B', 'C']
pers = itertools.permutations(iterable, 2)
print(list(pers))
"""
[
('A', 'B'), 
('A', 'C'), 
('B', 'A'), 
('B', 'C'), 
('C', 'A'), 
('C', 'B')
]
"""

 

product(iterable, repeat=1)

모든 조합을 포함하는 중복 수열을 생성 (중복된 요소의 선택도 허용)

import itertools

iterable = ['A','B','C']
products = itertools.product(iterable, repeat=2)

print(list(products))
'''
[
('A', 'A'), 
('A', 'B'), 
('A', 'C'), 
('B', 'A'), 
('B', 'B'), 
('B', 'C'), 
('C', 'A'), 
('C', 'B'), 
('C', 'C')
]
'''

 

따라서 중복 수열을 생성하는 product 메서드를 활용하여 아래와 같은 풀이가 가능하다.

import sys
data_temp = sys.stdin if sys.platform == 'linux' else open('입력.txt', 'r')
input_data = data_temp.read().splitlines()

from itertools import product 

def solution (data) :
  [_,m], iterable = map(lambda x : list(map(int, x.split())),data)
  products = list(product(iterable, repeat=m))
  result = list(map(lambda x : ' '.join(map(str, x)), products))
  print('\n'.join(result))
  
solution(input_data)

 

참고한 공식 문서 https://docs.python.org/3/library/itertools.html

itertools — Functions creating iterators for efficient looping