python itertools 사용 방법

itertools

itertools란 효율적인 루핑을 위한 이터레이터를 만드는 라이브러리로 여러가지 함수들이 많이 있지만,

그 중 코딩 테스트에서 대표적으로 사용되는 combinations(), combinations_with_replacement(), product(), permutations()에 대해 알아봅니다.

✅순열 (permutations)

permutations는 반복 가능한 객체(=길이가 n인)에 대해서 중복을 허용하지 않고 r개를 뽑아서 나열하는 함수입니다.

뽑힌 순서대로 나열하기 때문에 순서가 의미가 있습니다.

(즉, 같은 값이 뽑히더라도 순서가 다르면 다른 경우의 수로 취급합니다.)

만약 r을 지정하지 않거나 None으로 지정하면 최대 길이의 순열이 리턴 됩니다.

순열 예제 (1, 2, 3, 4 중에서 2개를 선택하여 만든 모든 순열):

• (1, 2), (1, 3), (1, 4)
• (2, 1), (2, 3), (2, 4)
• (3, 1), (3, 2), (3, 4)
• (4, 1), (4, 2), (4, 3)

🤔 어떤 문제로 나올 수 있을까?

순서가 중요한 경우에 사용됩니다.

예를 들어, 주어진 숫자들을 재배열하여 만들 수 있는 가장 큰 수나 가장 작은 수를 찾는 문제, 또는 문자열의 모든 가능한 순서를 찾는 문제 등이 있습니다.

예제 문제: 주어진 숫자들을 재배열하여 만들 수 있는 모든 순서의 숫자들 중 K번째로 큰 수를 찾는 문제.

itertools를 쓴 코드는 다음과 같습니다.

from itertools import permutations

s_list = ['A', 'B', 'C']

print(list(permutations(s_list)))

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

라이브러리를 사용하지 않은 코드는 다음과 같습니다.

def permutation(arr, r):
    result = []

    def generate(chosen, remaining):
        if len(chosen) == r:
            result.append(chosen[:])
            return
        for i in range(len(remaining)):
            chosen.append(remaining[i])
            generate(chosen, remaining[:i] + remaining[i+1:])
            chosen.pop()

    generate([], arr)
    return result

✅조합(combinations)

combinations는 반복 가능한 객체(=길이가 n인)에 대해서 중복을 허용하지 않고 r개를 뽑아 나열하는 함수입니다.

어떤 것을 뽑는지만 중요하게 보기 때문에 뽑은 순서는 고려하지 않습니다.

조합 예제 (1, 2, 3, 4 중에서 2개를 선택하는 모든 조합):

• (1, 2), (1, 3), (1, 4)
• (2, 3), (2, 4)
• (3, 4)

🤔 어떤 문제로 나올 수 있을까?

순서에 상관없이 특정 수의 요소를 선택할 때 사용됩니다.

예를 들어, 주어진 요소들 중 몇 개를 선택해 특정 조건을 만족하는 조합을 찾는 문제가 이에 해당합니다.

예제 문제: N개의 숫자 중에서 3개를 선택해 그 합이 특정 수 M에 가장 가까운 조합을 찾는 문제.

itertools를 쓴 코드는 다음과 같습니다.

from itertools import combinations
  for i in combinations([1,2,3,4], 2):
       print(i)

"""
(1, 2)
(1, 3)
(1, 4)
(2, 3)
(2, 4)
(3, 4)
"""

라이브러리를 사용하지 않은 코드는 다음과 같습니다.

def combination(arr, r):
    result = []

    def generate(chosen, start):
        if len(chosen) == r:
            result.append(chosen[:])
            return
        for i in range(start, len(arr)):
            chosen.append(arr[i])
            generate(chosen, i + 1)
            chosen.pop()

    generate([], 0)
    return result

✅중복 순열(product)

중복 가능한 n개에서 r개를 택하여 일렬로 나열하는 함수입니다.

중복순열은 순열과는 다르게 같은 숫자를 중복하여 사용할 수 있습니다.

중복순열 예제 (1, 2에서 중복을 허용하여 2자리 수 만드는 모든 경우):

• (1, 1), (1, 2)
• (2, 1), (2, 2)

🤔 어떤 문제로 나올 수 있을까?

같은 요소를 반복해서 선택할 수 있고, 순서가 중요한 경우에 사용됩니다.

예를 들어, 주어진 요소들로 만들 수 있는 모든 수열을 찾는 문제 등이 있습니다.

예제 문제: 주어진 숫자들을 사용하여 만들 수 있는 모든 R자리의 수열 중 특정 조건을 만족하는 경우의 수를 찾는 문제.

itertools를 쓴 코드는 다음과 같습니다.

from itertools import product

sets = [1,2,3]

#2개를 뽑아 일렬로 나열하는 경우의 수(단, 중복 허용)
data = product(sets, repeat = 2)

for i in data:
   print(i)

"""
(1, 1)
(1, 2)
(1, 3)
(2, 1)
(2, 2)
(2, 3)
(3, 1)
(3, 2)
(3, 3)
"""

라이브러리를 사용하지 않은 코드는 다음과 같습니다.

def product(arr, r):
    result = []

    def generate(chosen):
        if len(chosen) == r:
            result.append(chosen[:])
            return
        for i in range(len(arr)):
            chosen.append(arr[i])
            generate(chosen)
            chosen.pop()

    generate([])
    return result

✅중복 조합 (combinations_with_replacement)

중복 가능한 n개에서 순서를 생각하지 않고 r개를 택하여 나열하는 함수입니다.

중복조합 예제 (1, 2, 3에서 중복을 허용하여 2개를 선택하는 모든 경우):

• (1, 1), (1, 2), (1, 3)
• (2, 2), (2, 3)
• (3, 3)

🤔 어떤 문제로 나올 수 있을까?

같은 요소를 반복해서 선택할 수 있으며, 순서가 중요하지 않은 경우에 사용됩니다.

이는 주로 합이나 조합의 수를 찾는 문제에 활용됩니다.

예제 문제: 주어진 숫자들을 중복을 허용하여 선택해, 그 합이 특정 수가 되는 모든 조합을 찾는 문제.

itertools를 쓴 코드는 다음과 같습니다.

from itertools import combinations_with_replacement

for val in combinations_with_replacement([1,2,3,4], 2):
    print(val)
    
"""
(1, 1)
(1, 2)
(1, 3)
(1, 4)
(2, 2)
(2, 3)
(2, 4)
(3, 3)
(3, 4)
(4, 4)
"""

라이브러리를 사용하지 않은 코드는 다음과 같습니다.

def combinations_with_replacement(arr, r):
    # 중복조합을 저장할 리스트
    result = []

    # arr에서 r개를 중복 허용하여 선택하는 모든 경우를 찾는 함수
    def generate(chosen, start):
        if len(chosen) == r:
            result.append(chosen[:])
            return
        for i in range(start, len(arr)):
            chosen.append(arr[i])
            generate(chosen, i)
            chosen.pop()

    generate([], 0)
    return result