Blog

Python 중첩 리스트 평탄화(flatten) — 6가지 방법 성능 실측과 함정

API 응답을 페이지별로 모은 리스트, 파일마다 읽은 줄 목록처럼 리스트의 리스트를 한 줄로 펴야 하는 경우는 흔하다. 방법은 list comprehension, itertools.chain, sum(xss, []) 등 여럿인데, 어느 것이 빠른지에 대한 통념과 실제 측정값이 다르고, 깊이가 섞인 데이터에서는 일부 방법이 아예 동작하지 않는다.

이 글에서는 한 단계 평탄화 방법 여섯 가지를 Python 3.12.3에서 직접 측정해 비교하고, 임의 깊이 평탄화와 실제로 재현한 함정까지 정리한다.

비교한 방법과 측정 코드

import functools
import itertools
import operator
import timeit

METHODS = {
    "comprehension":   lambda xss: [x for xs in xss for x in xs],
    "extend loop":     lambda xss: _extend(xss),
    "chain.from_iter": lambda xss: list(itertools.chain.from_iterable(xss)),
    "chain(*xss)":     lambda xss: list(itertools.chain(*xss)),
    "reduce(iconcat)": lambda xss: functools.reduce(operator.iconcat, xss, []),
    "sum(xss, [])":    lambda xss: sum(xss, []),
}


def _extend(xss):
    out = []
    for xs in xss:
        out.extend(xs)
    return out


def bench(xss, number):
    for name, fn in METHODS.items():
        t = min(timeit.repeat(lambda: fn(xss), number=number, repeat=5)) / number
        print(f"  {name:<16} {t * 1e3:8.3f} ms")


xss = [[1, 2, 3], [4, 5, 6], [7], [8, 9]]
assert all(fn(xss) == [1, 2, 3, 4, 5, 6, 7, 8, 9] for fn in METHODS.values())

print("짧은 리스트 10,000개 (각 10개)")
bench([list(range(10)) for _ in range(10_000)], number=20)
print("긴 리스트 10개 (각 100,000개)")
bench([list(range(100_000)) for _ in range(10)], number=5)

측정값은 timeit.repeat 5회 중 최솟값을 1회 실행 시간으로 환산했다. 측정 환경은 Ubuntu 24.04 VirtualBox 게스트(6 vCPU)라 절대값보다 방법 간 비율을 보면 된다.

측정 결과

$ python3 flatten_bench.py
짧은 리스트 10,000개 (각 10개)
  comprehension       0.724 ms
  extend loop         0.207 ms
  chain.from_iter     0.538 ms
  chain(*xss)         0.556 ms
  reduce(iconcat)     0.207 ms
  sum(xss, [])      531.965 ms
긴 리스트 10개 (각 100,000개)
  comprehension       7.515 ms
  extend loop         2.896 ms
  chain.from_iter     5.042 ms
  chain(*xss)         4.378 ms
  reduce(iconcat)     2.699 ms
  sum(xss, [])       21.778 ms

흔히 가장 빠르다고 알려진 comprehension이 실제로는 extend보다 2.5~3.5배 느리다. extend와 iconcat은 원소를 하나씩 바이트코드로 옮기지 않고 C 레벨에서 리스트 단위로 복사하기 때문이다.

import timeit

for n in (1_000, 2_000, 4_000, 8_000):
    xss = [list(range(10)) for _ in range(n)]
    s = min(timeit.repeat(lambda: sum(xss, []), number=3, repeat=3)) / 3
    c = min(timeit.repeat(lambda: [x for xs in xss for x in xs], number=3, repeat=3)) / 3
    print(f"n={n:>5}  sum: {s * 1e3:8.2f} ms   comprehension: {c * 1e3:6.3f} ms")
$ python3 sum_scale.py
n= 1000  sum:     5.29 ms   comprehension:  0.107 ms
n= 2000  sum:    26.75 ms   comprehension:  0.248 ms
n= 4000  sum:   116.11 ms   comprehension:  0.277 ms
n= 8000  sum:   349.35 ms   comprehension:  0.861 ms

sum은 매 단계 a + b로 새 리스트를 만들어 지금까지의 결과를 통째로 복사하므로, 리스트 개수가 두 배가 될 때마다 시간이 3~5배씩 늘어난다(O(n²)).

방법짧은 리스트 다수긴 리스트 소수특징
[x for xs in xss for x in xs]0.72 ms7.5 ms가장 읽기 쉽고, 원소 변환·필터를 함께 넣을 수 있다
out.extend(xs) 루프0.21 ms2.9 ms빠르고 명확하다
chain.from_iterable(xss)0.54 ms5.0 ms지연 평가, 제너레이터 입력에 적합
chain(*xss)0.56 ms4.4 ms인자를 전부 풀어야 시작한다
reduce(iconcat, xss, [])0.21 ms2.7 ms가장 빠르지만 초기값을 빼면 원본을 바꾼다
sum(xss, [])532 ms21.8 msO(n²), 쓰지 말 것

깊이가 섞인 중첩 평탄화

위 방법은 모두 한 단계만 편다. [1, [2, [3]]]처럼 깊이가 제각각이면 재귀 제너레이터를 쓴다.

from collections.abc import Iterable

def deep_flatten(items):
    for x in items:
        if isinstance(x, Iterable) and not isinstance(x, (str, bytes)):
            yield from deep_flatten(x)
        else:
            yield x

print(list(deep_flatten([1, [2, [3, (4, 5)]], "abc", [[["xy"]]]])))
$ python3 deep.py
[1, 2, 3, 4, 5, 'abc', 'xy']

튜플도 함께 펴진다. 리스트만 펴고 싶으면 조건을 isinstance(x, list)로 좁힌다.

주의사항

아래 출력은 모두 같은 환경에서 실제로 실행한 결과다.

import functools
import itertools
import operator
from collections.abc import Iterable

# 1) reduce에 초기값을 빼면 첫 번째 리스트가 바뀐다
xss = [[1, 2], [3], [4, 5]]
flat = functools.reduce(operator.iconcat, xss)
print("1)", flat, "xss[0] =", xss[0], "same object:", flat is xss[0])

# 2) 재귀 평탄화에서 str 예외 처리를 빼면 무한 재귀
def naive(items):
    for x in items:
        if isinstance(x, Iterable):
            yield from naive(x)
        else:
            yield x

try:
    list(naive(["ab"]))
except RecursionError as e:
    print("2) RecursionError:", e)

# 3) 한 단계만 펴는 방법은 섞인 깊이에서 실패한다
try:
    print([x for xs in [1, [2, 3]] for x in xs])
except TypeError as e:
    print("3) TypeError:", e)

# 4) chain.from_iterable은 지연 평가, chain(*gen)은 전부 풀어야 시작
def rows():
    n = 0
    while True:
        n += 1
        yield [n, -n]

print("4)", list(itertools.islice(itertools.chain.from_iterable(rows()), 6)))
$ python3 pitfalls.py
1) [1, 2, 3, 4, 5] xss[0] = [1, 2, 3, 4, 5] same object: True
2) RecursionError: maximum recursion depth exceeded
3) TypeError: 'int' object is not iterable
4) [1, -1, 2, -2, 3, -3]
  • 1) reduce(iconcat, xss)처럼 초기값 []을 빼면 첫 번째 리스트에 나머지를 이어 붙여 원본 xss[0]이 바뀐다. 반드시 세 번째 인자로 새 리스트를 준다.
  • 2) 문자열은 한 글자짜리 문자열을 내놓는 iterable이라, 재귀 평탄화에서 str/bytes를 제외하지 않으면 끝없이 내려간다.
  • 3) 한 단계 평탄화 방법은 모든 원소가 iterable이라고 가정한다. 스칼라가 섞이면 TypeError가 나므로 재귀 버전을 쓴다.
  • 4) 무한·대용량 제너레이터에는 chain.from_iterable을 쓴다. chain(*gen)은 인자 언패킹 때문에 제너레이터를 끝까지 소비하려 해서 무한 제너레이터에서는 끝나지 않는다(같은 코드를 chain(*rows())로 바꾸면 timeout 5에 걸려 종료됨).
  • 재귀 버전은 중첩 깊이가 기본 재귀 한도(1000)에 가까우면 RecursionError가 난다. 그 정도로 깊은 데이터는 명시적 스택으로 바꾼다.

마무리

한 단계 평탄화는 가독성이면 comprehension, 속도면 extend 루프나 reduce(iconcat, xss, []), 제너레이터 입력이면 chain.from_iterable을 쓴다. sum(xss, [])은 짧아 보여도 O(n²)이라 쓰지 않는다.

참고

Tags: , , , , , , , , , , ,

답글 남기기