Python itertools 완전정복: 코딩 테스트에서 가장 자주 쓰이는 것

파이썬의 itertools 는 코딩 테스트에서 강력한 무기입니다.

특히 완전탐색, 브루트포스, 조합 문제에서 효율적인 풀이를 가능하게 합니다.

 

이글에서는 코테에서 가장 자주 사용되는 itertools 함수 6개를 실전 예제와 함께 소개합니다.

 

다음은 코테 내 사용 비중을 감안하여 정리해 본 표입니다.

특히, 이거 없으면 문제 못푸나? 라고 궁금하신 분을 위해서 그냥 직접 구현도 가능한지에 대한 언급도 하였습니다.

순위	함수명	사용 비중	대체 가능성	설명
1️⃣	permutations()	높음	직접 구현 가능하지만 실수 잦음	순열 : 순서 있는 경우의 수 탐색에 최적
2️⃣	combinations()	높음	직접 구현 가능	조합 : 조합 문제에 자주 등장
3️⃣	product()	중간	이중 루프 등으로 대체 가능	중복순열 : 중복 허용 순열에 유용
4️⃣	combinations_with_replacement()	낮음	중첩 루프 또는 재귀로 구현 가능	중복조합 : 특정 케이스에서만 활용
5️⃣	accumulate()	낮음	prefix sum 등으로 대체	누적합 : 간단한 누적 계산 시 유용
6️⃣	groupby()	낮음	dict + loop로 구현 가능	그룹화 : 정렬 후 그룹핑 문제에 사용

 

 

아래에서 순서, 중복이라는 의미는

뽑을때 순서가 있는지

뽑을때 중복해서 뽑을 수 있는지

헷갈릴 수 있으니 반복 숙달하시기 바랍니다.

 

1. permutations() - 순열 (순서 O, 중복 X)

from itertools import permutations

# 예시: [1, 2, 3] 중 2개를 뽑아 순서 있게 나열
for p in permutations([1, 2, 3], 2):
    print(p)
    
# (1, 2)
# (1, 3)
# (2, 1)
# (2, 3)
# (3, 1)
# (3, 2)

📌 활용 문제: 외판원 문제 (TSP), 숫자 나열, 경로 탐색

예: 비밀번호 가능한 조합 찾기, 미로 이동 경로 모든 경우의 수

---

2. combinations() - 조합 (순서 X, 중복 X)

from itertools import combinations

# 예시: [1, 2, 3] 중 2개를 순서 상관없이 뽑기
for c in combinations([1, 2, 3], 2):
    print(c)
    
# (1, 2)
# (1, 3)
# (2, 3)

📌 활용 문제: 부분집합, 메뉴 구성, 후보 조합

예: 치킨집 고르기, 감시 피하기, 백준 1759 암호 만들기

---

3. product() - 데카르트 곱 (중복 순열)

from itertools import product

# 예시: [1, 2]와 [3, 4]에서 각각 1개씩 뽑아 가능한 모든 조합
for p in product([1, 2], [3, 4]):
    print(p)

# (1, 3)
# (1, 4)
# (2, 3)
# (2, 4)

# 또는 같은 리스트 반복
for p in product([1, 2], repeat=2):
    print(p)
    
# (1, 1)
# (1, 2)
# (2, 1)
# (2, 2)

📌 활용 문제: 중복 가능한 숫자 고르기, 완전탐색, 사칙연산 시뮬레이션

예: 연산자 끼워넣기 (백준 14888)

---

4. combinations_with_replacement() - 중복 조합 (순서 X, 중복 O)

from itertools import combinations_with_replacement

# 예시: [1, 2]에서 중복 허용해 2개 조합
for c in combinations_with_replacement([1, 2], 2):
    print(c)
    
# (1, 1)
# (1, 2)
# (2, 2)

📌 활용 문제: 중복 동전 조합, 무한 공급 아이템 구성

예: 동전 2 (백준 2294), 중복 집합 구성

---

5. accumulate() - 누적합 계산

from itertools import accumulate

# 예시: 누적합 계산
from operator import mul

print(list(accumulate([1, 2, 3, 4])))        # [1, 3, 6, 10]
print(list(accumulate([1, 2, 3, 4], mul)))  # [1, 2, 6, 24]

📌 활용 문제: 구간합, 최댓값 누적, DP 보조 계산

예: 누적 거리 계산, 연속된 수 계산

---

6. groupby() - 연속된 값 그룹화

from itertools import groupby

# 예시: 같은 값이 연속된 그룹 묶기
for key, group in groupby('AAABBBCCDAA'):
    print(key, list(group))
    
# A ['A', 'A', 'A']
# B ['B', 'B', 'B']
# C ['C', 'C']
# D ['D']
# A ['A', 'A']

📌 활용 문제: 정렬 후 카운팅, 연속된 수 묶기

예: 압축, 연속된 방문자 그룹 세기