Itertools - chain

Python으로 알고리즘 공부를 하다보면 리스트(list)를 다뤄야할 상황이 빈번히 생깁니다. 물론 데이터분석이나 머신러닝, 딥러닝을 위해선 Numpy나 Pandas 라이브러리를 사용할 수 있지만, 코딩테스트를 위해 공부를 하다보면 내장 라이브러리외 사용이 금지되어있습니다.

 

itertools에서 제공되는 함수들은 내장 라이브러리기 때문에 알아두면 유용한데, 오늘 소개할 함수는 chain입니다.

그럼 이 함수를 언제 쓰느냐? 리스트를 붙여 하나의 리스트로 만들 때 쓰게 됩니다.

 

"어? 리스트는 단순히 리스트끼리 덧셈(+)을 쓰면 되지 않나요?" 할 수 도있습니다. 당연히 리스트끼리 더하면 더욱 편리하지만 리스트 배열이 2차원일 경우는 조금 복잡하고 시간복잡도를 계산하면서 코드를 짜기 위해선 알아야하는 함수입니다.

 

우선 리스트끼리의 덧셈을 하게 되면 어떻게 되는지 확인해보면 다음과 같습니다.

a = ['지구온난화', '삼림', '감소']
b = ['원숭이', '영장류', '생존']

print(a+b)
# ['지구온난화', '삼림', '감소', '원숭이', '영장류', '생존']

단순히 2개의 리스트만 존재했을 땐 이렇지만 붙여야할 리스트의 갯수를 모를 때도 있고 리스트 명을 모를 때도 있고 다양한 이유로 힘들 가능성이 있기 때문에 단순히 loop문을 돌려 해결할 수 없습니다. 물론 다음과 같은 방법으로 구할 순 있습니다.

list_names = [a, b, c, ..., z]
total = []
for x in list_names:
	total.append(x)

벌써 코드가 깔끔하지 않으면서 조금 귀찮아지는 현상이 발생합니다. 이럴 때 itertools에서 제공하는 chain을 사용하면 깔끔하게 작성을 하며 속도 또한 줄일 수 있습니다.

from itertools import chain

a = ['지구온난화', '삼림', '감소']
b = ['원숭이', '영장류', '생존']

print(list(chain(a, b)))

이렇게 되면 처음 단순히 더하는 것보다 길어지겠지만 저렇게 리스트를 나눠서 저장하는 것이 아니라 2차원 배열로 저장하면서 효율적으로 구성을 하게 되면 더더욱 좋습니다.

from itertools import chain

a = [['지구온난화', '삼림', '감소'],['원숭이', '영장류', '생존']]

print(list(chain.from_iterable(a)))
# ['지구온난화', '삼림', '감소', '원숭이', '영장류', '생존']

속도 면에선 chain을 사용하면 약 2배가 빠른걸로 측정이 된다.

import math
import time

start = time.time()
list(set(chain.from_iterable(train_tokenized)))
end = time.time()

print(f"{end - start:.5f} sec")
#0.00009 sec

start = time.time()
tmp = []
for x in train_tokenized:
  for t in x:
    tmp.append(t)
list(set(tmp))
end = time.time()

print(f"{end - start:.5f} sec")
#0.00021 sec