백준 N과 M 시리즈 itertools로 정복하기
in Algorithm on Algorithm, Python
백준 N과 M 시리즈 문제를 Python 내장 모듈인 itertools로 풀어보자.
개요
Python의 내장 모듈 itertools는 다양한 반복자를 제공한다.
이를 이용하면 반복문을 사용하지 않고도 풀 수 있는 문제들이 많다.
특히 ‘N과 M’ 시리즈의 문제는 수열을 다양한 방식으로 출력해야 하는 문제로 itertools를 사용하면 매우 간단하게 풀 수 있다.
이번 글에서는 백준 ‘N과 M’ 시리즈의 문제를 Python 내장 모듈인 itertools로 풀어보자.
itertools
본격적으로 문제를 풀기 전에 간략하게 itertools에 대해 알아보자.
- 조합형 반복자: 이번 포스트에서 자주 쓰일 함수들이다.
product(*iterables, repeat=1):iterables의 카테시안 곱을 생성한다.from itertools import product print(list(product([1, 2, 3], [4, 5, 6]))) # [# 4 5 6 # 1 (1, 4), (1, 5), (1, 6), # 2 (2, 4), (2, 5), (2, 6), # 3 (3, 4), (3, 5), (3, 6), # ] print(list(product([1, 2], [3, 4], [5, 6]))) # [ 1 5 6 # 3 (1, 3, 5), (1, 3, 6), # 4 (1, 4, 5), (1, 4, 6), # 2 # 3 (2, 3, 5), (2, 3, 6), # 4 (2, 4, 5), (2, 4, 6) # ]permutations(iterable, r=None):iterable의 길이가r인 순열을 생성한다.from itertools import permutations print(list(permutations([1, 2, 3]))) # [(1, 2, 3), (1, 3, 2), (2, 1, 3), (2, 3, 1), (3, 1, 2), (3, 2, 1)] # permutations(iterable) == permutations(iterable, len(iterable)) print(list(permutations([1, 2, 3], 2))) # [(1, 2), (1, 3), (2, 1), (2, 3), (3, 1), (3, 2)]combinations(iterable, r):iterable의 길이가r인 조합을 생성한다. 원소 별 중복을 허용하지 않는다.from itertools import combinations print(list(combinations([1, 2, 3], 2))) # [(1, 2), (1, 3), (2, 3)]combinations_with_replacement(iterable, r):iterable의 길이가r인 조합을 생성한다. 원소 별 중복을 허용한다.from itertools import combinations_with_replacement print(list(combinations_with_replacement([1, 2, 3], 2))) # [(1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)]
- 그외 유한 반복자
accumulate(iterable[, func, *, initial=None]):iterable의 누적값 배열을 생성한다.func가 주어지지 않으면operator.add를 사용한다.initial이 주어지면iterable의 첫번째 원소 앞에initial을 추가한다. 최종 누적값만이 필요한 경우iterable.reduce를 사용하는 것이 더 빠르다.from itertools import accumulate from operator import add, mul print(list(accumulate([1, 2, 3, 4, 5]))) # [1, 3, 6, 10, 15] # accumulate([1, 2, 3, 4, 5]) == accumulate([1, 2, 3, 4, 5], add) = accumulate([1, 2, 3, 4, 5], lambda x, y: x + y) print(list(accumulate([1, 2, 3, 4, 5], mul))) # [1, 2, 6, 24, 120] print(list(accumulate([1, 2, 3, 4, 5], mul, initial=100))) # [100, 100, 200, 600, 2400, 12000]chain(*iterables):iterables의 원소를 연결하여 하나의 반복자로 만든다.from itertools import chain print(list(chain([1, 2, 3], [4, 5, 6]))) # [1, 2, 3, 4, 5, 6]iterables가 하나의iterable로 구성된 경우,chain.from_iterable(iterable)를 이용하면 된다.from itertools import chain print(list(chain.from_iterable([[1, 2, 3], [4, 5, 6]]))) # [1, 2, 3, 4, 5, 6]compress(data, selectors):selectors가truly한 요소와 동일한 인덱스의data요소를 생성한다.from itertools import compress print(list(compress(range(10), [0, 0, 1, 1, 0, 1, 0, 1, 0, 0]))) # [2, 3, 5, 7]dropwhile(predicate, iterable):predicate가truly한 요소를 만날 때까지iterable을 건너뛴다.itertools.takewhile의 반대이다.from itertools import dropwhile print(list(dropwhile(lambda x: x < 5, [1, 4, 6, 4, 1]))) # [6, 4, 1]filterfalse(predicate, iterable):predicate가falsy한 요소를 생성한다. 내장함수인filter와 반대의 경우이다.from itertools import filterfalse print(list(filterfalse(lambda x: x % 2, range(10)))) # [0, 2, 4, 6, 8]groupby(iterable, key=None):iterable의 연속된 같은 요소를 그룹으로 묶어서 생성한다.key가 주어지면key의 반환값이 같은 요소를 그룹으로 묶는다.from itertools import groupby print(list((key, tuple(group)) for key, group in groupby('AAAABBBCCDAABBB'))) # [ # ('A', ('A', 'A', 'A', 'A')), # ('B', ('B', 'B', 'B')), # ('C', ('C', 'C')), # ('D', ('D',)), # ('A', ('A', 'A')), # ('B', ('B', 'B', 'B')), # ] print(list((key, tuple(group)) for key, group in groupby('AaaBBbcCAAa', str.lower))) # [ # ('a', ('A', 'a', 'a')), # ('b', ('B', 'B', 'b')), # ('c', ('c', 'C')), # ('a', ('A', 'A', 'a')), # ]연속되지 않은 같은 키의 그룹끼리 묶고 싶다면
groupby를 사용하기 전에sorted를 사용하자.print(list((key, tuple(group)) for key, group in groupby(sorted('AaaBBbcCAAa', key=str.lower), str.lower))) # [ # ('a', ('A', 'a', 'a', 'A', 'A', 'a')), # ('b', ('B', 'B', 'b')), # ('c', ('c', 'C')), # ]islice(iterable, stop),islice(iterable, start, stop[, step]):iterable의stop번째 요소까지 생성한다.start가 주어지면start번째 요소부터 생성한다.step이 주어지면step만큼 건너뛴 요소를 생성한다.stop이None이면iterable가 끝날 때까지 생성한다.Sequence의slice와 유사하다.from itertools import islice print(list(islice('ABCDEFG', 2))) # ['A', 'B'] print(list(islice('ABCDEFG', 2, 4))) # ['C', 'D'] print(list(islice('ABCDEFG', 2, None))) # ['C', 'D', 'E', 'F', 'G'] print(list(islice('ABCDEFG', 1, 5, 2))) # ['B', 'D'] print(list(islice('ABCDEFG', 1, None, 2))) # ['B', 'D', 'F']pairwise(iterable):iterable의 인접한 두 요소를 튜플로 묶어서 생성한다.from itertools import pairwise print(list(pairwise('ABCDEFG'))) # [('A', 'B'), ('B', 'C'), ('C', 'D'), ('D', 'E'), ('E', 'F'), ('F', 'G')]starmap(function, iterable):iterable의 요소를function에 넘겨서 생성한다.map(function, *iterable)과 유사하다.from itertools import starmap print(list(starmap(pow, [(2, 5), (3, 2), (10, 3)]))) # [32, 9, 1000]takewhile(predicate, iterable):predicate가truly인 동안iterable의 요소를 생성한다.dropwhile의 반대이다.from itertools import takewhile print(list(takewhile(lambda x: x < 5, [1, 4, 6, 4, 1]))) # [1, 4]tee(iterable, n=2):iterable을n개의 복사본으로 분리한다.iterable을 여러 번 사용해야 할 때 유용하다.from itertools import tee a, b, c = tee('ABC', 3) print(list(a)) # ['A', 'B', 'C'] print(list(b)) # ['A', 'B', 'C'] print(list(c)) # ['A', 'B', 'C']zip_longest(*iterables, fillvalue=None):iterables의 요소를 튜플로 묶어서 생성한다.zip과 유사하지만,zip은 최단iterable이 끝날 때까지만 생성하는 데에 반해zip_longest는 최장iterable이 끝날 때까지 생성한다. 최장iterable보다 짧은iterable의 요소는fillvalue로 채운다.from itertools import zip_longest print(list(zip_longest('ABCD', 'xy', fillvalue='-'))) # [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')] # zip 이었다면 [('A', 'x'), ('B', 'y')]
- 무한 반복자
count(start=0, step=1):start부터step씩 증가하는 수열을 생성한다.from itertools import count, islice print(list(islice(count(2, 3), 5, 20, 2))) # [17, 23, 29, 35, 41, 47, 53, 59]cycle(iterable):iterable을 무한 반복한다.from itertools import cycle, islice;print(list(islice(cycle('ABC'), 7))) # ['A', 'B', 'C', 'A', 'B', 'C', 'A']repeat(object, times=None):object를times번 반복한다.times가 생략되면 무한 반복한다.from itertools import repeat, islice print(list(islice(repeat(7), 5))) # [7, 7, 7, 7, 7]
문제 풀이
본격적으로 문제를 풀어보자.
N과 M(1)
1부터 N까지 자연수 중에서 중복 없이 M개를 고른 수열을 사전순으로 출력해야한다. 중복 없는 수열이므로 순열 즉, itertools.permutations를 사용한다.
from itertools import permutations
N, M = map(int, input().split())
nums = [str(i) for i in range(1, N + 1)]
# join을 사용하기 위해 미리 문자열로 변환
combs = permutations(nums, M)
print('\n'.join(' '.join(comb) for comb in combs))
N과 M(2)
1부터 N까지 자연수 중에서 중복 없이 M개를 고른 오름차순의 수열을 사전순으로 출력해야한다. 중복 없는 오름차순 수열이므로 조합 즉, itertools.combinations를 사용한다.
from itertools import combinations
N, M = map(int, input().split())
nums = [str(i) for i in range(1, N + 1)]
combs = combinations(nums, M)
print('\n'.join(' '.join(comb) for comb in combs))
N과 M(3)
1부터 N까지 자연수 중에서 M개를 고른 수열을 사전순으로 출력해야한다. 중복이 허용되므로 itertools.product를 사용한다.
from itertools import product
N, M = map(int, input().split())
nums = [str(i) for i in range(1, N + 1)]
combs = product(nums, repeat=M)
print('\n'.join(' '.join(comb) for comb in combs))
N과 M(4)
1부터 N까지 자연수 중에서 M개를 고른 비내림차순 수열을 사전순으로 출력해야한다. 비내림차순 수열이므로 중복을 허용하는 조합 즉, itertools.combinations_with_replacement를 사용한다.
from itertools import combinations_with_replacement
N, M = map(int, input().split())
nums = [str(i) for i in range(1, N + 1)]
combs = combinations_with_replacement(nums, M)
print('\n'.join(' '.join(comb) for comb in combs))
N과 M(5)
주어진 중복되지 않는 N개의 수 중에서 중복 없이 M개를 고른 수열을 사전순으로 출력해야한다. 1번과 비슷하지만 주어진 수열이 있으므로 이를 정렬해야 한다.
from itertools import permutations
N, M = map(int, input().split())
nums = sorted(input().split(), key=int)
# key를 int로 주지 않으면 문자열로 정렬되므로 주의
# sorted(["10", "2"]) -> ["10", "2"]
# sorted(["10", "2"], key=int) -> ["2", "10"]
combs = permutations(nums, M)
print('\n'.join(' '.join(comb) for comb in combs))
N과 M(6)
주어진 중복되지 않는 N개의 수 중에서 중복 없이 M개를 고른 오름차순 수열을 사전순으로 출력해야한다. 2번과 비슷하지만 주어진 수열이 있으므로 이를 정렬해야 한다.
from itertools import combinations
N, M = map(int, input().split())
nums = sorted(input().split(), key=int)
combs = combinations(nums, M)
print('\n'.join(' '.join(comb) for comb in combs))
N과 M(7)
주어진 중복되지 않는 N개의 수 중에서 M개를 고른 오름차순 수열을 사전순으로 출력해야한다. 3번과 비슷하지만 주어진 수열이 있으므로 이를 정렬해야 한다.
from itertools import product
N, M = map(int, input().split())
nums = sorted(input().split(), key=int)
combs = product(nums, repeat=M)
print('\n'.join(' '.join(comb) for comb in combs))
N과 M(8)
주어진 중복되지 않는 N개의 수 중에서 M개를 고른 비내림차순 수열을 사전순으로 출력해야한다. 4번과 비슷하지만 주어진 수열이 있으므로 이를 정렬해야 한다.
from itertools import combinations_with_replacement
N, M = map(int, input().split())
nums = sorted(input().split(), key=int)
combs = combinations_with_replacement(nums, M)
print('\n'.join(' '.join(comb) for comb in combs))
N과 M(9)
주어진 N개의 수 중에서 중복 없이 M개를 고른 수열을 사전순으로 출력해야한다. 5번과 비슷하지만 중복된 수가 주어질 수 있으므로 수열을 정렬해야 한다.
from itertools import permutations
N, M = map(int, input().split())
nums = map(int, input().split())
combs = sorted(set(permutations(nums, M)))
print('\n'.join(' '.join(map(str, comb)) for comb in combs))
N과 M(10)
주어진 N개의 수 중에서 중복 없이 M개를 고른 비내림차순 수열을 사전순으로 출력해야한다. 8번과 비슷하지만 중복된 수가 주어질 수 있으므로 주어진 수와 수열을 정렬해야 한다.
from itertools import combinations
N, M = map(int, input().split())
nums = sorted(map(int, input().split()))
combs = sorted(set(combinations(nums, M)))
print('\n'.join(' '.join(map(str, comb)) for comb in combs))
N과 M(11)
주어진 N개의 수 중에서 중복 없이 M개를 고른 수열을 사전순으로 출력해야한다. 7번과 비슷하지만 중복된 수가 주어질 수 있으므로 주어진 수를 정렬해야 한다.
from itertools import product
N, M = map(int, input().split())
nums = sorted(set(map(int, input().split())), key=int)
combs = product(nums, repeat=M)
print('\n'.join(' '.join(map(str, comb)) for comb in combs))
N과 M(12)
주어진 N개의 수 중에서 M개를 고른 비내림차순 수열을 사전순으로 출력해야한다. 9번과 비슷하지만 중복된 수가 주어질 수 있으므로 주어진 수와 수열을 정렬해야 한다.
from itertools import combinations_with_replacement
N, M = map(int, input().split())
nums = sorted(set(map(int, input().split())), key=int)
combs = sorted(set(combinations_with_replacement(nums, M)))
print('\n'.join(' '.join(map(str, comb)) for comb in combs))
여담
비슷한 이름의 N과 M이라는 문제가 있다. 하지만 전혀 다른 문제이므로 여기에 풀이를 적진 않을 것이다.