Table of Contents
  1. Iterators terminating on the shortest input sequence
    1. accumulate 🔗
    2. compress 🔗
    3. dropwhile 🔗
    4. takewhile 🔗
    5. groupby 🔗
    6. zip_longest 🔗
  2. Combinatoric iterators
    1. product 🔗
    2. permutations 🔗
    3. combinations 🔗
    4. combinations_with_replacement 🔗
  3. 마치며

이번 포스팅에서는 Python 2.7과 3.5부터 탑재된 itertools의 유용하고 자주 쓸만한 함수들을 익혀보도록 하겠습니다.

사실 for와 if만 알고 있다면 굳이 Itertools를 쓰지 않고도 구현 가능한 기능들이지만, 고급 개발자들이 고심하여 만든 코드이므로 한번 보는 것만으로도 도움이 되리라 생각합니다.

itertools를 다음과 같이 import하여, 이후 코드에서는 it라고 선언하겠습니다.
1
import itertools as it

Iterators terminating on the shortest input sequence

accumulate 🔗

배열을 받아, 누적 합으로 이루어진 배열을 반환합니다.
(기본적으로 sum을 하나, 별도의 operator function을 인자로 줄 수 있습니다.)

python3.8에서는 initial 인자가 추가되었습니다.

python 버전 별 accumulate 함수
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
# python3.8
def accumulate(iterable, func=operator.add, *, initial=None):
    it = iter(iterable)
    total = initial
    if initial is None:
        try:
            total = next(it)
        except StopIteration:
            return
    yield total
    for element in it:
        total = func(total, element)
        yield total

# python3.7 이하
def accumulate(iterable, func=operator.add):
    it = iter(iterable)
    try:
        total = next(it)
    except StopIteration:
        return
    yield total
    for element in it:
        total = func(total, element)
        yield total
자연수 배열 예제입니다.
1
2
3
4
5
6
7
8
9
10
for i in it.accumulate([1, 2, 3, 4, 5]):
    print(i)

"""
1  # 1
3  # 1 + 2
6  # 1 + 2 + 3
10 # 1 + 2 + 3 + 4
15 # 1 + 2 + 3 + 4 + 5
"""
숫자 뿐만 아니라 문자열도 연산 가능합니다.
1
2
3
4
5
6
7
8
9
10
for i in it.accumulate(['a', 'b', 'c', 'd', 'e']):
    print(i)

"""
a
ab
abc
abcd
abcde
"""

다른 함수를 인자로 주어 다양한 연산도 가능합니다.

operator의 곱하기 합수를 주었습니다.
1
2
3
4
5
6
7
8
9
10
11
import operator as op
for i in it.accumulate([1,2,3,4,5], op.mul):
    print(i)

"""
1   # 1
2   # 1 * 2
6   # 1 * 2 * 3
24  # 1 * 2 * 3 * 4
120 # 1 * 2 * 3 * 4 * 5
"""
max 함수를 이용하여 다음과 같은 연산도 가능합니다.
1
2
3
4
5
6
7
8
9
10
for i in it.accumulate([1,2,8,3,7], max):
    print(i)

"""
1
2
8
8
8
"""

compress 🔗

문자열과 True/False 배열을 받아서, True인 인덱스의 문자열 값만 반환합니다.


python3.8 compress function
1
2
def compress(data, selectors):
    return (d for d, s in zip(data, selectors) if s)
배열에서 첫 번째와 세 번째 인자만 True이므로 a와 c만 반환했습니다.
1
2
3
4
5
6
7
for i in it.compress('abcd', [1, 0, 1, 0]):
    print(i)

"""
a
b
"""

dropwhile 🔗

function과 iterable을 받아서 False 처음 등장하기 전까지는 모두 버리고, 나머지를 반환합니다.

python3.8 dropwhile function
1
2
3
4
5
6
7
8
def dropwhile(predicate, iterable):
    iterable = iter(iterable)
    for x in iterable:
        if not predicate(x):
            yield x
            break
    for x in iterable:
        yield x
5번째 인자인 6이 나오기 전까지는 True를 반환하므로 모두 버리고 6과 4만 출력되었습니다.
1
2
3
4
5
6
7
for i in it.dropwhile(lambda x: x<5, [1, 4, 3, 4, 6, 4]):
    print(i)

"""
6
4
"""

takewhile 🔗

dropwhile과 반대로 처음 False가 등장할 때까지만 반환합니다.

python3.8 takewhile function
1
2
3
4
5
6
def takewhile(predicate, iterable):
    for x in iterable:
        if predicate(x):
            yield x
        else:
            break
5번째 인자인 6이 나오기 전까지만 True를 반환하므로 6부터 나머지는 반환하지 않습니다.
1
2
3
4
5
6
7
8
9
for i in it.takewhile(lambda x: x<5, [1, 4, 3, 4, 6, 4]):
    print(i)

"""
1
4
3
4
"""

groupby 🔗

문자열을 입력받아 연속하여 등장하는 문자열을 그룹으로 묶어서 반환합니다.

python3.8 groupby class
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class groupby:
    def __init__(self, iterable, key=None):
        if key is None:
            key = lambda x: x
        self.keyfunc = key
        self.it = iter(iterable)
        self.tgtkey = self.currkey = self.currvalue = object()
    def __iter__(self):
        return self
    def __next__(self):
        self.id = object()
        while self.currkey == self.tgtkey:
            self.currvalue = next(self.it)    # Exit on StopIteration
            self.currkey = self.keyfunc(self.currvalue)
        self.tgtkey = self.currkey
        return (self.currkey, self._grouper(self.tgtkey, self.id))
    def _grouper(self, tgtkey, id):
        while self.id is id and self.currkey == tgtkey:
            yield self.currvalue
            try:
                self.currvalue = next(self.it)
            except StopIteration:
                return
            self.currkey = self.keyfunc(self.currvalue)
groupby()는 각 그룹의 문자열과 해당 문자열이 반복된 리스트를 반환합니다.
1
2
3
4
5
6
7
8
9
10
11
for k, g in it.groupby('AAAABBBCCDAABBB'):
    print(k, list(g))

"""
A ['A', 'A', 'A', 'A']
B ['B', 'B', 'B']
C ['C', 'C']
D ['D']
A ['A', 'A']
B ['B', 'B', 'B']
"""

zip_longest 🔗

내장함수 zip과 달리 가장 긴 인자를 기준으로 zip 연산을 합니다.

python3.8 zip_longest function
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
def zip_longest(*args, fillvalue=None):
    iterators = [iter(it) for it in args]
    num_active = len(iterators)
    if not num_active:
        return
    while True:
        values = []
        for i, it in enumerate(iterators):
            try:
                value = next(it)
            except StopIteration:
                num_active -= 1
                if not num_active:
                    return
                iterators[i] = repeat(fillvalue)
                value = fillvalue
            values.append(value)
        yield tuple(values)

zip_longest를 예제에 앞서 내장함수 zip에 대해 알아보겠습니다.

zip은 길이가 짧은 인자를 기준으로 요소들을 결합한 배열을 반환합니다.

'abc'가 'ㄱㄴㄷㄹㅁ' 보다 길이가 짧으므로 'abc' 기준으로 결과가 반환 된 것을 볼 수 있습니다.
1
2
3
4
5
list(zip('abc', 'ㄱㄴㄷㄹㅁ'))

"""
[('a', 'ㄱ'), ('b', 'ㄴ'), ('c', 'ㄷ')]
"""
zip_longest는 길이가 긴 'ㄱㄴㄷㄹㅁ'를 기준으로 모자란 부분은 None을 채워서 반환하는 것을 볼 수 있습니다.
1
2
3
4
5
list(it.zip_longest('abc', 'ㄱㄴㄷㄹㅁ'))

"""
[('a', 'ㄱ'), ('b', 'ㄴ'), ('c', 'ㄷ'), (None, 'ㄹ'), (None, 'ㅁ')]
"""

Combinatoric iterators

product 🔗

product는 iterable의 곱집합을 구하는 함수입니다.

python3.8 product function
1
2
3
4
5
6
7
def product(*args, repeat=1):
    pools = [tuple(pool) for pool in args] * repeat
    result = [[]]
    for pool in pools:
        result = [x+[y] for x in result for y in pool]
    for prod in result:
        yield tuple(prod)
'ABCD'와 'xy'의 곱집합인 'Ax Ay Bx By Cx Cy Dx Dy'를 반환하는 예제입니다.
1
2
3
4
5
6
7
8
9
10
11
12
list(it.product('ABCD', 'xy'))

"""
[('A', 'x'),
 ('A', 'y'),
 ('B', 'x'),
 ('B', 'y'),
 ('C', 'x'),
 ('C', 'y'),
 ('D', 'x'),
 ('D', 'y')]
"""
repeat 인자를 넣어주어 다음이 활용할 수도 있습니다.
1
2
3
4
5
6
7
8
9
10
11
12
list(it.product(range(2), repeat=3))

"""
[(0, 0, 0),
 (0, 0, 1),
 (0, 1, 0),
 (0, 1, 1),
 (1, 0, 0),
 (1, 0, 1),
 (1, 1, 0),
 (1, 1, 1)]
"""

1과 -1로된 크기가 3인 배열의 경우의 수는 다음과 같습니다.

1
2
3
4
5
6
7
8
[(1, 1, 1),
 (1, 1, -1),
 (1, -1, 1),
 (1, -1, -1),
 (-1, 1, 1),
 (-1, 1, -1),
 (-1, -1, 1),
 (-1, -1, -1)]

위 함수를 product를 이용하면 한줄로 만들 수 있습니다.

1
list(it.product([1, -1], repeat=3))

permutations 🔗

순열을 뜻하는 permutations은 n개 중에서 r개를 뽑아 만들 수 있는 순서있는 나열입니다.

python3.8 permutations function
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
def permutations(iterable, r=None):
    pool = tuple(iterable)
    n = len(pool)
    r = n if r is None else r
    if r > n:
        return
    indices = list(range(n))
    cycles = list(range(n, n-r, -1))
    yield tuple(pool[i] for i in indices[:r])
    while n:
        for i in reversed(range(r)):
            cycles[i] -= 1
            if cycles[i] == 0:
                indices[i:] = indices[i+1:] + indices[i:i+1]
                cycles[i] = n - i
            else:
                j = cycles[i]
                indices[i], indices[-j] = indices[-j], indices[i]
                yield tuple(pool[i] for i in indices[:r])
                break
        else:
            return

순열의 개수는 다음과 같이 구할 수 있습니다.

nPr = n!/(n-r)!

두 번째 인자는 iterable에서 뽑는 개수이므로 len(iterable) 보다 클 수 없습니다.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
list(it.permutations('ABCD', 2))

"""
[('A', 'B'),
 ('A', 'C'),
 ('A', 'D'),
 ('B', 'A'),
 ('B', 'C'),
 ('B', 'D'),
 ('C', 'A'),
 ('C', 'B'),
 ('C', 'D'),
 ('D', 'A'),
 ('D', 'B'),
 ('D', 'C')]
"""
r을 생략할 경우 iterable 크기만큼 뽑아서 만들 수 있는 순열을 반환합니다.
1
2
3
4
5
list(it.permutations(range(3)))

"""
[(0, 1, 2), (0, 2, 1), (1, 0, 2), (1, 2, 0), (2, 0, 1), (2, 1, 0)]
"""

combinations 🔗

조합을 의미하는 combinations는 n개 중에서 r를 뽑아 만들 수 있는 집합(set)의 나열입니다. (한 번 뽑은 인자는 다시 뽑을 수 없습니다.)

python3.8 combinations function
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
def combinations(iterable, r):
    pool = tuple(iterable)
    n = len(pool)
    if r > n:
        return
    indices = list(range(r))
    yield tuple(pool[i] for i in indices)
    while True:
        for i in reversed(range(r)):
            if indices[i] != i + n - r:
                break
        else:
            return
        indices[i] += 1
        for j in range(i+1, r):
            indices[j] = indices[j-1] + 1
        yield tuple(pool[i] for i in indices)

조합의 개수는 다음과 같이 구할 수 있습니다.

nCr = nPr/r! = n!/(r!*(n-r)!)

'ABCD' 네 개중에서 순서를 무시하고 2개를 조합하는 방법은 6가지(4C2) 입니다.
1
2
3
4
5
list(it.combinations('ABCD', 2))

"""
[('A', 'B'), ('A', 'C'), ('A', 'D'), ('B', 'C'), ('B', 'D'), ('C', 'D')]
"""

combinations_with_replacement 🔗

위에서 설명한 combinations과 달리 한 번 뽑은 인자를 다시 뽑을 수 있는 경우의 조합입니다.

python3.8 combinations_with_replacement
1
2
3
4
5
6
7
8
9
10
11
12
13
14
pool = tuple(iterable)
n = len(pool)
if not n and r:
    return
indices = [0] * r
yield tuple(pool[i] for i in indices)
while True:
    for i in reversed(range(r)):
        if indices[i] != n - 1:
            break
    else:
        return
    indices[i:] = [indices[i] + 1] * (r - i)
    yield tuple(pool[i] for i in indices)
'A'를 뽑은 뒤 'A'를 또 뽑을 수 있습니다.
1
2
3
4
5
list(it.combinations_with_replacement('ABC', 2))

"""
[('A', 'A'), ('A', 'B'), ('A', 'C'), ('B', 'B'), ('B', 'C'), ('C', 'C')]
"""

마치며

성능에 최적화된 기본 내장기능인 Itertools에 대해 알아보았습니다. 처음에 언급하였듯 직접 만들 수 있는 기능들이지만, 바퀴를 다시 발명 할 필요는 없다고 생각합니다. 본 포스팅에서 설명하지 않은 함수들은 공식문서를 참고하시기 바랍니다.