[Python advanced] itertools.groupby() 쉬운 설명 및 다양한 예제

파이썬 기본 라이브러리인 itertools에는 groupby라는 함수가 있다. 알고리즘 문제풀 때나 개발 할 때, groupby를 활용한 사고가 익숙해지기만 하면 정말 유용하게 쓰일 수 있어서 개인적으로 좋아하는 함수다.

 

본 글에서는 itertools패키지의 groupby가 무엇인지 설명하고, 이를 활용한 몇 가지 알고리즘 문제를 풀어보려고 한다.

 

itertools란?

매개변수로 입력된 iterable안에 있는 키를 확인 후, 분류하여 동일한 키를 가진 element 들을 해당 키를 기준으로 그룹지어 리턴한다.

 

 

itertools.groupby 예제

이런건 사실 예제로 봐야 이해가 된다.

 

from itertools import groupby

d = [('Europe','Manchester'),
     ('America','NewYork'),
     ('Asia','Seoul'),
     ('Asia','Tokyo'),
     ('America','Chicago'),
     ('America','Seattle'),
     ('Europe','London'),
     ('Asia','Beijing'),
     ('Europe','Paris'),
     ]

# Key function
key_func = lambda x: x[0]

g_result = groupby(d, key_func)
for key, group in g_result:
    print(key + " :", list(group))

위 예시를 출력해보면,

Europe : [('Europe', 'Manchester')]
America : [('America', 'NewYork')]
Asia : [('Asia', 'Seoul'), ('Asia', 'Tokyo')]
America : [('America', 'Chicago'), ('America', 'Seattle')]
Europe : [('Europe', 'London')]
Asia : [('Asia', 'Beijing')]
Europe : [('Europe', 'Paris')]

groupby 다양한 예시

이걸 어디에 활용해볼 수 있을까? 한번 예시들로 보자

 

 

리스트의 특정 값을 기준으로 묶기

my_list = [["A",0], ["B",1], ["C",0], ["D",2], ["E",2]]

위와 같은 list가 있다고 했을 때, 두 번째 값을 기준으로 모으고 싶다면 어떻게 해야 할까.

아래처럼 groupby와 정렬을 활용해 쉽게 풀 수 있다.

from itertools import groupby

my_list = [["A",0], ["B",1], ["C",0], ["D",2], ["E",2]]
my_list.sort(key= lambda x: x[1])
# Key function
key_func = lambda x: x[1]

g_result = groupby(my_list, key_func)
for key, group in g_result:
    print(list(group)[0])

 

 

홀수는 홀수끼리, 짝수는 짝수끼리 모으기

my_list = list(range(0,11))

위와 같은 list가 있다고 했을 때, 두 번째 값을 기준으로 모으고 싶다면 어떻게 해야 할까.

먼저 아래 예시처럼 하게되면, 출력이 이상하게 되는데, 그 이유는 groupby는 consecutive한 element 들 끼리만 group지어주기 때문이다.

from itertools import groupby

my_list = list(range(0,11))

# Key function
key_func = lambda x: "even" if x%2 ==0 else "odd"

g_result = groupby(my_list, key_func)
for key, group in g_result:
    # print(list(group)[0])
    print(str(key) + " :", list(group))


# print(dict(g_result))

그래서 이렇게 하려면 위에서 정렬한 것처럼 똑같이 정렬을 해주어야 한다.

from itertools import groupby
import random 

my_list = list(range(0,11))

# Key function
key_func = lambda x: "even" if x%2 ==0 else "odd"

#shuffled to verify sorting only considers key_func
random.shuffle(my_list)
print(my_list)
#sort my_list by Key function before groupby
my_list.sort(key = key_func)

print("sorted my_list", my_list)

g_result = groupby(my_list, key_func)
for key, group in g_result:
    # print(list(group)[0])
    print(str(key) + " :", list(group))

 

 

A, B 리스트에 있을 때, 같은 number로 correspond하는 A element들을 묶는 법

A = [5, 7, 9, 12, 8, 16, 25]
B = [2, 1, 3, 2, 3, 1, 4]

위 처럼 있을 때, 같은 인덱스에 있는 값들이 같은 것들 끼리 묶는다면?

[7,16] corresponding to the number 1 of listB
[5, 12] corresponding to the number 2 of listB
[9, 8] corresponding to the number 3 of listB
[25] corresponding to the number 4 of listB

위처럼 될텐데, 이걸 어떻게 할 수 있을까?

아래처럼 groupby를 활용하면 쉽게 풀 수 있다.

from itertools import groupby
import random 

A = [5, 7, 9, 12, 8, 16, 25]
B = [2, 1, 3, 2, 3, 1, 4]

#zip elements
zipped = list(zip(A, B))
print(zipped)

key_func = lambda x: x[1]
zipped.sort(key = key_func)
print("after sort", zipped)

g_result = groupby(zipped, key_func)
for key, group in g_result:
    print(key, ":", list(group))

 

 

리스트 속 연속된 수 끼리 group짓는 법

given the set {0, 1, 2, 3, 4, 7, 8, 9, 11} I want to get { {0,4}, {7,9}, {11,11} }.

from itertools import groupby
l = [0, 1, 2, 3, 4, 7, 8, 9, 11]

# build tuples from the list items and their respective index
# (index, value)
e = list(enumerate(l))
print(e)

#group tuples using difference between their index and their value
key_func = lambda x : x[1]-x[0]
g_result = groupby(e, key_func)

#print result
# for key, group in g_result:
#     print(key, ":", list(group))

grouped = []
#groupthem
for key, group in g_result:
    g_list = list(group)
    grouped.append((g_list[0][1], g_list[-1][1]))

print(grouped)

 

 

같은 prefix를 가지는 longest sequence 찾기

potential_labels = ['foo', 'foo::bar', 'foo::bar::baz', "abc", "abc::cde::def", "bleh"]

위와 같은 label이 있을때, 아래처럼 첫번째 foo, abc와 같은 string을 가지는데, 제일 긴 string을 구하려면 어떻게 해야 할까?

['foo::bar::baz', 'abc::cde::def', 'bleh']

from itertools import groupby
potential_labels = ['foo', 'foo::bar', 'foo::bar::baz', "abc", "abc::cde::def", "bleh"]

key_func = lambda x: x.split(":")[0]

g_result = groupby(potential_labels, key_func)

#print groupby result
# for k, group in g_result:
#     print(k, ":", list(group))


#get max length result
print([max(group, key=len) for _, group in g_result])

 

참고자료

 

http://pyengine.blogspot.com/2018/01/groupby.html

https://www.geeksforgeeks.org/itertools-groupby-in-python/

 

위에서 풀어본 문제도 포함이고, 다양한 문제 예시는 아래를 참고하면 된다.

• Anagrams (see notebook)
• Binning
• Group odd and even numbers
• Group a list by values
• Remove duplicate elements
• Find indices of repeated elements in an array
• Split an array into n-sized chunks
• Find corresponding elements between two lists
• Compression algorithm (see notebook)/Run Length Encoding
• Grouping letters by length, key function (see notebook)
• Consecutive values over a threshold (see notebook)
• Find ranges of numbers in a list or continuous items (see docs)
• Find all related longest sequences
• Take consecutive sequences that meet a condition (see related post)