파이썬 itertools로 이터레이터 마음껏 조합하기

파이썬으로 데이터를 다루다 보면 여러 리스트를 하나로 이어 붙이거나, 조건을 만족할 때까지만 순회하거나, 모든 조합을 구해야 하는 상황이 생기곤 하는데요. 그럴 때마다 for 루프와 임시 리스트를 조합해 작성하다 보면 코드가 생각보다 길어집니다. 😅

이럴 때 파이썬 표준 라이브러리의 itertools 모듈이 큰 도움이 됩니다. 이터레이터를 효율적으로 만들고 조합하는 함수를 모아놓은 모듈인데요. 내장 함수인 map()이나 filter()와 마찬가지로 결과를 한꺼번에 메모리에 올리지 않고 필요할 때마다 하나씩 생성하기 때문에 대용량 데이터를 다룰 때도 효율적이에요.

이번 포스팅에서는 itertools에서 실무에서 가장 자주 쓰이는 함수를 실용적인 예제와 함께 살펴보겠습니다.

itertools 모듈 불러오기

itertools는 파이썬 표준 라이브러리에 포함되어 있어서 별도로 설치할 필요가 없습니다. 사용하려는 함수만 골라서 임포트하거나 모듈 전체를 가져오면 됩니다.

from itertools import chain, islice, product
# 또는
import itertools

chain — 여러 이터러블을 하나로 이어 붙이기

여러 리스트를 하나로 합쳐서 순회해야 할 때 가장 먼저 떠올릴 함수입니다. + 연산자로 리스트를 합치면 새 리스트가 메모리에 생성되지만, chain()은 이터레이터를 반환하기 때문에 메모리 면에서 유리합니다.

from itertools import chain

fruits = ["사과", "배", "포도"]
veggies = ["당근", "브로콜리", "시금치"]
grains = ["쌀", "보리", "밀"]

for item in chain(fruits, veggies, grains):
    print(item)

사과
배
포도
당근
브로콜리
시금치
쌀
보리
밀

중첩 리스트(리스트의 리스트)를 평탄화(flatten)할 때는 chain.from_iterable()이 더 편리합니다.

from itertools import chain

nested = [["사과", "배"], ["당근", "브로콜리"], ["쌀", "보리"]]

print(list(chain.from_iterable(nested)))

['사과', '배', '당근', '브로콜리', '쌀', '보리']

islice — 이터레이터를 슬라이싱하기

일반 리스트는 my_list[2:5]처럼 슬라이싱할 수 있지만, 이터레이터는 인덱싱을 지원하지 않습니다. islice()는 이터레이터에도 슬라이싱과 유사한 기능을 제공합니다.

from itertools import islice

data = range(100)

# 처음 5개
print(list(islice(data, 5)))

# 인덱스 10부터 15까지 (start, stop)
print(list(islice(data, 10, 15)))

# 인덱스 0부터 20까지, 2칸 간격 (start, stop, step)
print(list(islice(data, 0, 20, 2)))

[0, 1, 2, 3, 4]
[10, 11, 12, 13, 14]
[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]

로그 파일이나 API 응답처럼 길이를 미리 알 수 없는 이터레이터에서 앞부분만 미리 보고 싶을 때 특히 유용합니다.

takewhile / dropwhile — 조건 기반으로 자르기

takewhile()은 조건 함수가 True인 동안만 요소를 가져오고, 조건이 처음으로 False가 되는 순간 멈춥니다. dropwhile()은 반대로 조건이 True인 동안 요소를 건너뛰다가 처음 False가 되는 시점부터 나머지를 모두 반환합니다.

from itertools import takewhile, dropwhile

temps = [18, 22, 25, 28, 30, 27, 22, 18]

# 30도 미만인 동안만 가져오기
print(list(takewhile(lambda t: t < 30, temps)))

# 30도 미만인 동안 건너뛰기
print(list(dropwhile(lambda t: t < 30, temps)))

[18, 22, 25, 28]
[30, 27, 22, 18]

정렬된 데이터를 다룰 때 특히 빛을 발합니다. 예를 들어 타임스탬프 순으로 정렬된 로그에서 특정 시간 이후의 항목만 가져오거나, 특정 시간 이전의 항목만 처리할 때 활용할 수 있습니다.

filterfalse — 조건을 만족하지 않는 요소 추출하기

파이썬 내장 filter()는 조건 함수가 True인 요소만 반환하는데요. filterfalse()는 반대로 조건 함수가 False인 요소만 반환합니다.

from itertools import filterfalse

numbers = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

# 홀수만 (짝수가 아닌 것)
odds = list(filterfalse(lambda n: n % 2 == 0, numbers))
print(odds)

[1, 3, 5, 7, 9]

filter(lambda n: n % 2 != 0, numbers)와 동일하지만, 조건을 부정하지 않아도 되니 코드 의도가 더 명확하게 드러날 때가 있습니다.

starmap — 언패킹이 필요한 함수에 map 적용하기

파이썬 map() 함수는 map(func, iterable) 형태로 사용하는데, 인자가 여러 개인 함수에는 직접 사용하기 어렵습니다. starmap()은 각 요소를 튜플로 받아 언패킹(*)하면서 함수를 호출해 줍니다.

from itertools import starmap
import operator

pairs = [(2, 3), (4, 5), (10, 2)]

# operator.pow(base, exp) → base ** exp
print(list(starmap(operator.pow, pairs)))

# 같은 결과를 map으로 쓰려면 언패킹이 필요해서 더 복잡해짐
print(list(map(lambda p: p[0] ** p[1], pairs)))

[8, 1024, 100]
[8, 1024, 100]

데이터베이스 쿼리 결과나 CSV 파일처럼 행(row)이 여러 컬럼으로 구성된 데이터를 처리할 때 편리합니다.

groupby — 연속된 요소를 키로 묶기

groupby()는 연속으로 이어진 같은 키를 가진 요소를 묶어줍니다. 데이터베이스의 GROUP BY와 이름이 같지만, 반드시 같은 키를 가진 요소가 연속으로 배치되어 있어야 한다는 점이 다릅니다. 그래서 보통 sorted()로 먼저 정렬한 뒤에 씁니다.

from itertools import groupby

products = [
    {"name": "사과", "category": "과일"},
    {"name": "배", "category": "과일"},
    {"name": "당근", "category": "채소"},
    {"name": "시금치", "category": "채소"},
    {"name": "포도", "category": "과일"},
]

# category 기준으로 정렬 후 그룹핑
sorted_products = sorted(products, key=lambda p: p["category"])

for category, items in groupby(sorted_products, key=lambda p: p["category"]):
    names = [item["name"] for item in items]
    print(f"{category}: {names}")

과일: ['사과', '배', '포도']
채소: ['당근', '시금치']

조합 이터레이터 — product, permutations, combinations

itertools의 꽃이라고 할 수 있는 조합 관련 함수들입니다. 알고리즘 문제나 경우의 수를 구해야 하는 상황에서 반복문을 여러 번 중첩하지 않아도 됩니다.

product — 카르테시안 곱 (중복 순열)

두 이터러블의 모든 조합을 구합니다. 중첩 for 루프를 대체합니다.

from itertools import product

colors = ["빨강", "파랑"]
sizes = ["S", "M", "L"]

for color, size in product(colors, sizes):
    print(f"{color}-{size}")

빨강-S
빨강-M
빨강-L
파랑-S
파랑-M
파랑-L

같은 이터러블 내에서 중복을 허용한 경우의 수를 구하려면 repeat 인자를 사용합니다. 예를 들어 동전을 3번 던졌을 때 모든 경우의 수는 다음과 같습니다.

from itertools import product

for result in product(["앞", "뒤"], repeat=3):
    print(result)

('앞', '앞', '앞')
('앞', '앞', '뒤')
('앞', '뒤', '앞')
('앞', '뒤', '뒤')
('뒤', '앞', '앞')
('뒤', '앞', '뒤')
('뒤', '뒤', '앞')
('뒤', '뒤', '뒤')

permutations — 순열 (순서가 중요한 배열)

주어진 요소들로 만들 수 있는 모든 순열을 반환합니다. 두 번째 인자로 길이를 지정할 수 있습니다.

from itertools import permutations

# 3명 중 1등, 2등, 3등을 뽑는 경우의 수
runners = ["Alice", "Bob", "Charlie"]

for perm in permutations(runners):
    print(perm)

('Alice', 'Bob', 'Charlie')
('Alice', 'Charlie', 'Bob')
('Bob', 'Alice', 'Charlie')
('Bob', 'Charlie', 'Alice')
('Charlie', 'Alice', 'Bob')
('Charlie', 'Bob', 'Alice')

4명 중 2명을 순서 있게 뽑는다면 이렇게 씁니다.

from itertools import permutations

runners = ["Alice", "Bob", "Charlie", "Dave"]

results = list(permutations(runners, 2))
print(f"경우의 수: {len(results)}")
print(results[:5])

경우의 수: 12
[('Alice', 'Bob'), ('Alice', 'Charlie'), ('Alice', 'Dave'), ('Bob', 'Alice'), ('Bob', 'Charlie')]

combinations — 조합 (순서가 중요하지 않은 선택)

순열과 달리 순서를 고려하지 않은 조합을 반환합니다. ('Alice', 'Bob')과 ('Bob', 'Alice')를 같은 것으로 취급합니다.

from itertools import combinations

# 4명 중 2명을 팀으로 묶는 경우의 수
players = ["Alice", "Bob", "Charlie", "Dave"]

results = list(combinations(players, 2))
print(f"경우의 수: {len(results)}")
for combo in results:
    print(combo)

경우의 수: 6
('Alice', 'Bob')
('Alice', 'Charlie')
('Alice', 'Dave')
('Bob', 'Charlie')
('Bob', 'Dave')
('Charlie', 'Dave')

중복을 허용한 조합이 필요하다면 combinations_with_replacement()를 사용합니다.

from itertools import combinations_with_replacement

# 3가지 메뉴 중 2개를 고를 때 같은 메뉴 2개도 허용
menus = ["피자", "파스타", "샐러드"]

for combo in combinations_with_replacement(menus, 2):
    print(combo)

('피자', '피자')
('피자', '파스타')
('피자', '샐러드')
('파스타', '파스타')
('파스타', '샐러드')
('샐러드', '샐러드')

무한 이터레이터 — count, cycle, repeat

itertools에는 무한히 값을 생성하는 이터레이터도 있습니다. 무한 루프처럼 들리지만 islice()나 takewhile()과 함께 쓰면 안전하게 활용할 수 있습니다.

count — 숫자를 계속 세기

from itertools import count, islice

# 10부터 2씩 증가하는 숫자 5개
print(list(islice(count(10, 2), 5)))

[10, 12, 14, 16, 18]

zip()과 함께 사용하면 enumerate()처럼 인덱스를 붙일 수도 있습니다. 단, 시작 번호나 간격을 자유롭게 지정할 수 있다는 점이 다릅니다.

from itertools import count

items = ["사과", "배", "포도"]
for idx, item in zip(count(1), items):
    print(f"{idx}. {item}")

1. 사과
2. 배
3. 포도

cycle — 이터러블을 계속 반복하기

from itertools import cycle, islice

# 3가지 패턴을 7번 반복
pattern = cycle(["A", "B", "C"])
print(list(islice(pattern, 7)))

['A', 'B', 'C', 'A', 'B', 'C', 'A']

repeat — 같은 값을 반복하기

from itertools import repeat

# 같은 값 5번 반복
print(list(repeat("hello", 5)))

# map()의 두 번째 인자처럼 활용
print(list(map(pow, range(5), repeat(2))))

['hello', 'hello', 'hello', 'hello', 'hello']
[0, 1, 4, 9, 16]

accumulate — 누적 계산하기

accumulate()는 functools.reduce()와 비슷하지만, 최종 결과 하나만 반환하는 reduce()와 달리 중간 누적값을 모두 반환합니다.

from itertools import accumulate
import operator

numbers = [1, 2, 3, 4, 5]

# 기본값: 누적 합계
print(list(accumulate(numbers)))

# 누적 곱
print(list(accumulate(numbers, operator.mul)))

# 누적 최댓값
print(list(accumulate(numbers, max)))

[1, 3, 6, 10, 15]
[1, 2, 6, 24, 120]
[1, 2, 3, 4, 5]

주가 데이터의 누적 최고가나 누적 합계 그래프를 그릴 때처럼 중간 결과가 모두 필요한 상황에서 유용합니다.

마치며

지금까지 itertools 모듈에서 자주 사용하는 함수를 살펴봤습니다. 이 모듈이 강력한 이유는 단순히 편의 함수를 모아놓은 것 이상으로, 모든 함수가 이터레이터를 반환해서 필요할 때만 값을 생성하는 지연 평가(lazy evaluation) 방식으로 동작하기 때문이에요. 덕분에 대용량 데이터를 처리할 때도 메모리를 효율적으로 사용할 수 있습니다.

정리하면 이런 상황에서 itertools를 떠올리면 좋습니다.

• chain() — 여러 이터러블을 하나로 연결
• islice() — 이터레이터를 슬라이싱
• takewhile() / dropwhile() — 조건 기반으로 자르거나 건너뛰기
• groupby() — 연속된 요소를 키로 그룹핑
• product(), permutations(), combinations() — 카르테시안 곱, 순열, 조합
• accumulate() — 누적 계산 중간 결과 반환

itertools와 함께 활용하면 좋은 파이썬의 map() 함수는 파이썬 map() 함수 포스팅에서, filter() 함수는 파이썬 filter() 함수 포스팅에서 자세히 다루고 있으니 함께 읽어 보세요.

더 깊이 파고들고 싶다면 파이썬 공식 itertools 문서를 추천합니다. 레시피 섹션에 실전에서 바로 쓸 수 있는 패턴이 많이 정리되어 있거든요.