python itertools

파이썬 표준 라이브러리 중 하나로, 반복해서 작성해야하는 코드를 구현할 수 있는 라이브러리이다.
데카르트 곱(product), 순열(permutations), 조합(combinations), 중복조합(combinations_with_replacement)을 구현하기에 편리하다.

product

'데카르트 곱' 또는 'cartesian product'를 표현할 때 사용하는 메소드이다.
'product(A, B)'은 '((x,y) for x in A for y in B)'과 표현하는 메소드이다.

product 정의 내용

def product(*args, repeat=1):
    # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy
    # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111
    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)

product(A, repeat=4)은 product(A, A, A, A)와 같다.

product 사용 예시

from itertools import product

l = [(1,-1), (2, -2), (3, -3)]
a = list(product(*l))
print(a)
# [(1, 2, 3), (1, 2, -3), (1, -2, 3), (1, -2, -3), (-1, 2, 3), (-1, 2, -3), (-1, -2, 3), (-1, -2, -3)]

b = list(product('ABCD', repeat=2))
print(b)
# [AA, AB, AC, AD, BA, BB, BC, BD, CA, CB, CC, CD, DA, DB, DC, DD]

permutations

순열을 표현(permutations(n, r))할 때 사용하는 메소드이다. r을 생략하면 n으로 대체된다.

permutations 정의 내용

def permutations(iterable, r=None):
    # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
    # permutations(range(3)) --> 012 021 102 120 201 210
    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

permutations 사용 예시

from itertools import permutations

l = [1,2,3]
a = list(permutations(l))
print(a)
# [(1, 2, 3), (1, 3, 2), (2, 1, 3), (2, 3, 1), (3, 1, 2), (3, 2, 1)]

b = list(permutations('ABCD', 2))
print(b)
# [AB, AC, AD, BA, BC, BD, CA, CB, CD, DA, DB, DC]

combinations

조합을 표현(combinations(n, r))할 때 사용하는 메소드이다. 이때 꼭 r을 넣어줘야 한다.

combinations 정의 내용

def combinations(iterable, r):
    # combinations('ABCD', 2) --> AB AC AD BC BD CD
    # combinations(range(4), 3) --> 012 013 023 123
    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)

combinations 사용 예시

from itertools import combinations

l = [1,2,3]
a = list(combinations(l, 2))
print(a)
# [(1, 2), (1, 3), (2, 3)]

b = list(combinations('ABCD', 2))
print(b)
# [AB, AC, AD, BC, BD, CD]

combinations_with_replacement

중복조합을 표현(combinations_with_replacement(n, r))할 때 사용하는 메소드이다. 이때 꼭 r을 넣어줘야 한다.

combinations_with_replacement 정의 내용

def combinations_with_replacement(iterable, r):
    pool = tuple(iterable)
    n = len(pool)
    for indices in product(range(n), repeat=r):
        if sorted(indices) == list(indices):
            yield tuple(pool[i] for i in indices)

combinations_with_replacement 사용 예시

from itertools import combinations_with_replacement

l = [1,2,3]
a = list(combinations_with_replacement(l, 2))
print(a)
# [(1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)]

b = list(combinations_with_replacement('ABCD', 2))
print(b)
# [AA, AB, AC, AD, BB, BC, BD, CC, CD, DD]

참고사이트

https://docs.python.org/3/library/itertools.html#itertools.permutations