[Python] iter함수와 itertools를 이용하여 메모리를 절약하면서 배열을 순서대로 그룹화(Grouping)하기

---

이번 포스팅에서는 itertools를 이용하여 배열을 같은 개수를 갖는 그룹으로 쪼개는 함수를 만들려고 한다.

 

예를 들어 배열(리스트) a = [ 1, 2, 3, 4]가 있다고 할 때 순서대로 2개 원소를 갖는 그룹을 만든다면 결과는 다음과 같다.

---

[(1, 2), (3, 4)]

---

이제 개념은 알았으니 이를 파이썬으로 구현하려면 어떻게 해야 할까?

---

   1. 원시적인 방법

다음과 같은 함수를 생각해볼 수 있다. 배열의 그룹 개수를 구하고 각 그룹 인덱스 별로 배열의 위치를 중복되지 않게 슬라이싱하는 방법으로 가져올 수 있다.

 

def naive_chunk(arr, n):
    num_groups = len(arr) // n
    return [tuple(arr[i*n:(i+1)*n]) for i in range(num_groups)]

 

c = [1,2,3,4]

print(naive_grouper(c, 2))

 

원시적인 그룹핑 결과

---

   2. iter함수를 활용한 그룹핑

이번엔 iter함수를 사용해보려고 한다. iter 함수는 배열을 iterator형식으로 만들어준다. iterator 형식은 리스트(list)나 튜플(tuple) 그리고 for 루프를 돌 때 비로소 배열 속 원소에 접근할 수 있으며 모든 원소를 다 접근했을 경우 없어지게 된다.

 

아래 코드를 실행하면 좀 더 쉽게 이해가 가능하다. 먼저 a는 iterator형식으로 암묵적으로 1, 2, 3, 4라는 원소를 갖고 있다. 하지만 for 루프를 만나서 원소에 접근하면 해당 원소는 없어지게 된다. for 루프 한번 돌고 break을 걸어주고 나서 a안에 남아 있는 원소를 보면 2, 3, 4만 남아 있는 것을 알 수 있다.

 

a = iter((1,2,3,4))
for x in a:
    print(x)
    break
print(list(a))

 

 

이러한 원리를 이용하여 원시적인 방법보다 더 세련된 방식으로 그룹핑 함수를 만들 수 있다.

 

def better_chunk(arr, n):
    iters = [iter(arr)] * n
    return zip(*iters)

c = [1,2,3,4]

print(list(better_chunk(c, 2)))

 

iter를 활용한 세련된 그룹핑 결과

어떻게 이것이 가능한가? 먼저 iters 변수에 동일한 iterator를 한 그룹이 갖는 원소 개수인 n개만큼 생성해준다. 동일한 iterator는 같은 메모리 주소를 가진다.

 

zip함수는 iter변수를 차례대로 돌면서 n=2개씩 튜플로 만드는데 이 과정이 정말 대박이다. 아래 그림은 zip에서 첫 번째 튜플이 만들어지는 과정이다. 첫 번째 iterator의 첫 번째 원소인 1을 zip의 첫 번째 튜플의 첫번째 원소로 가져온다. 첫 번째가 왜케 많아!! 헷갈린다..

 

그리고 zip함수가 두 번째 iterator의 첫 번째 원소를 가져와야하는데 1은 앞에서 사용했으므로 없어져서 2를 첫번째 튜플의 두 번째 원소로 가져오게 된다. 이런식으로 두 번째 튜플에는 1, 2 이가 iterator에서 사려졌으므로 최종적으로 3, 4를 가지고 앞의 과정을 반복하게 된다.

이 함수는 완전한 것이 아니다.

 

다음의 상황을 보자. 4개짜리의 배열을 3개씩 그룹핑하는 상황이다.

 

c = [1,2,3,4]

print(list(better_chunk(c, 3)))

 

 

1, 2, 3은 제대로 나왔으나 마지막 원소 4가 없어졌다. 왜냐하면 zip이 튜플을 만드는 과정에서 모든 원소를 다 소모해버렸기 때문에 중간에서 멈춰버린 것이기 때문이다(앞의 설명을 이해했다면 알 수 있다). 만약 4까지 나오게 하고 싶다면 이때 zip 대신에 itertools.zip_longest를 써보자. 가상의 원소(None)가 있는 것처럼 그룹핑을 해준다.

 

import itertools

def chunk(arr, n):
    iters = [iter(arr)] * n
    return itertools.zip_longest(*iters)

c = [1,2,3,4]

print(list(chunk(c, 3)))

 

 

만약 None을 없애고 싶다면 다음과 같이 바꿔준다.

 

def chunk(arr, n):
    iters = [iter(arr)] * n
    res = list(itertools.zip_longest(*iters))
    return (tuple(x for x in y if x is not None) for y in res)

print(list(chunk(c, 3)))

 

 

원시적인 방법보다 iter를 사용하는 것이 더 좋은 이유는 큰 배열을 다루는 데 있어서 메모리 절약 효과가 매우 크며 속도도 원시적인 방법보다 5배 정도 더 빠르다고 한다(참고 사이트)

 

파이썬에는 내가 모르는 참 신기한 것들이 많다. 가끔씩 찾아보고 공부해야겠다.

---

참고자료

https://realpython.com/python-itertools/

 

---