백준 N과 M 시리즈 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이라는 문제가 있다. 하지만 전혀 다른 문제이므로 여기에 풀이를 적진 않을 것이다.

Chomu's Blog.

Posts