19 przydatnych funkcji biblioteki Python Itertools (przykłady)

We wszystkim, co robimy, istnieją pewne powtarzalne problemy (oraz ich rozwiązania). W bibliotece standardowej Python Itertools znajdziemy zestaw wielu bardzo przydatnych funkcji tworzących iteratory. Wszystko z troski o to, żebyśmy nie wynajdowali koła na nowo i po prostu tworzyli kod szybciej 🙂

Wbudowane iteratory w Python Itertools

Biblioteka standardowa Python słynie z tego, że jest olbrzymia. Naprawdę wiele rzeczy jesteśmy w stanie zrobić bez zaciągania zewnętrznych zależności. To niewątpliwie jest duży plus. W tym wpisie opisałem, czym są iteratory i jakie niosą korzyści. A w tym przedstawię Ci, czym jest Python Itertools. Czyli zestaw 19. funkcji tworzących ciekawe iteratory.

Podział iteratorów w bibliotece itertools

W oficjalnej dokumentacji znajdziemy opis każdego iteratora z osobna wraz z jego poglądową implementacją. Dodatkowo iteratory są pogrupowane w trzech kategoriach: nieskończone iteratory, iteratory kończące się na najkrótszej sekwencji wejściowej oraz iteratory kombinatoryczne.

Przyjrzyjmy się im z bliska 🙂

Iteratory nieskończone

1. count

Tworzy iterator, który zwraca wartości, począwszy od tej podanej jako parametr start. Kolejne są zwiększane przez wartość step – drugi parametr.

Przykładowe użycie:

count_iterator = count()
for i in range(3):
    print(next(count_iterator))
# 0
# 1
# 2

Możemy zmienić wartość początkową:

count_iterator = count(1)
for i in range(3):
    print(next(count_iterator))
# 1
# 2
# 3

Również możemy zmienić wartość, o jaką zmieniają się kolejne liczby. W przypadku powyżej zwracamy tylko liczby parzyste:

count_iterator = count(step=2)
for i in range(3):
    print(next(count_iterator))
# 0
# 2
# 4

Oczywiście możemy ustawić ujemny krok:

count_iterator = count(1, step=-1)
for i in range(3):
    print(next(count_iterator))
# 1
# 0
# -1

Przy użyciu count możemy też odtworzyć działanie funkcji wbudowanej enumerate:

colors = ['red', 'blue', 'green']
indexed_colors = [color for color in zip(count(), colors)]
print(indexed_colors)
# [(0, 'red'), (1, 'blue'), (2, 'green')]

2. cycle

cycle przyjmuje jako parametr iterable. Ten iterator przechodzi przez przekazany argument iterable. Z tą różnicą, że jeśli się wyczerpie (dojdzie do końca) to nadal będzie zwracał wartości od początku. Dzieje się tak, bo cycle zapisuje sobie każdy element w pamięci.

cycle_iterator = cycle('kasjer')
for i in range(8):
    print(next(cycle_iterator))

# k
# a
# s
# j
# e
# r
# k
# a

Kolejny przykład to dzielenie uczestników na X grup przy użyciu zasady „Kolejno odlicz do X”.

users = ['Bob', 'Roman', 'Laszlo', 'Vlad', 'Ben', 'Greg', 'Leo', 'Adam']
cycle_iterator = cycle([1, 2, 3, 4])
users_teams = [name for name in zip(cycle_iterator, users)]
print(users_teams)
# [(1, 'Bob'), (2, 'Roman'), (3, 'Laszlo'), (4, 'Vlad'), (1, 'Ben'), (2, 'Greg'), (3, 'Leo'), (4, 'Adam')]

3. Repeat

Funkcja ta zwraca dany obiekt X razy. Jeśli parametr times nie jest przekazany, to obiekt będzie zwracany nieskończenie długo. Oficjalna dokumentacja tłumaczy, że repeat używa się głównie jako stały argument przekazywany do funkcji map. Może też służyć do dodawania stałej wartości do krotki (ang. tuple).

Poniższy przykład pochodzi z oficjalnej dokumentacji i oblicza potęgi kwadratowe dla cyfr z przedziału 0-9.

squares = list(map(pow, range(10), repeat(2)))
print(squares) # [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

Iteratory kończące się na najkrótszej sekwencji wejściowej

4. accumulate

Jest to funkcja tworząca zgromadzone sumy lub rezultaty funkcji binarnych. Zachowuje się podobnie do znanej funkcji reduce z biblioteki functools, z tą różnicą, że zwraca bieżące wartości dla każdego elementu przekazanego do funkcji.

Prosty przykład dodający do siebie kolejne elementy iterable:

numbers = [1, 2, 3, 4, 5]
result = list(accumulate(numbers))
print(result)  # [1, 3, 6, 10, 15]

Domyślnie wartości są do siebie dodawane. Jeśli chcemy zmienić operację, to można to zrobić, przekazując funkcję jako parametr func. Należy pamiętać, aby przyjmowała dokładnie dwa argumenty (tzw. funkcja binarna). Na przykład niech będzie to iloczyn:

numbers = [1, 2, 3, 4, 5]
result = list(accumulate(numbers, func=operator.mul))
print(result)  # [1, 2, 6, 24, 120]

Kolejnym przykładem może być zwracanie bieżącej wartości maksymalnej:

numbers = [1, 2, 1, 4, 3]
result = list(accumulate(numbers, func=max))
print(result)  # [1, 2, 2, 4, 4]

Ilość elementów w zwróconym iteratorze będzie równa ilości elementów w przekazanym jako argument iterable. Jedynym odstępstwem jest przekazanie opcjonalnego argumentu initial, który określa wartość początkową.

numbers = [1, 2, 1, 4, 3]
result = list(accumulate(numbers, func=max, initial=3))
print(result)  # [3, 3, 3, 3, 4, 4]

5. Chain

Bardzo przydatna funkcjonalność do iterowania po wielu zbiorach jeden po drugim. Najpierw iteruje po pierwszym iterable, następnie przechodzi do drugiego itd.

colors = ('red', 'orange')
numbers = (1, 2)
for element in chain(colors, numbers):
    print(element)
# red
# orange - koniec pierwszego iterable
# 1 - początek drugiego iterable
# 2

6. chain.from_iterable

Powyżej opisana funkcja chain posiada metodę from_iterable, której chyba najpopularniejsze zastosowanie, to wypłaszczenie zbioru.

input_list = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
result = chain.from_iterable(input_list)
print(list(result))  # [1, 2, 3, 4, 5, 6, 7, 8, 9]

7. compress

Przyjmuje jako parametr dwa iterowalne zbiory. Iterator zwrócony przez compress zwraca tylko wartości z pierwszego zbioru, dla których wartość w drugim zbiorze (selectors) na tej samej pozycji oznacza logiczną prawdę (True).

data = ['h', 'b', 'e', 'o', 'b', 'p']
selectors = [0, 1, 0, 1, 1, 0]
result = compress(data, selectors)
print(''.join(list(result)))  # bob

Gdy zbiór data jest większy niż zbiór selectors, to nadmiarowe elementy z automatu mają wartości fałszywe (False) – nie pojawią się w wynikowym iteratorze.

8. dropwhile

Zwraca iterator, który ignoruje/porzuca wszystkie wartości ze zbioru iterable do momentu gdy pierwszy element nie spełni warunku (parametr predicate). Od tej chwili każdy kolejny element będzie zwracany nawet jeśli nie spełnia on warunku, ponieważ nie jest on już sprawdzany.

input = [1, 2, 4, 6, 2, 1]
result = dropwhile(lambda x: x < 3, input)
print(list(result))  # [4, 6, 2, 1]

9. takewhile

Zachowuje się dokładnie odwrotnie do opisanego wyżej dropwhile. Zwraca wartości, dopóki warunek jest spełniony. Od pierwszego złamania warunku (parametr predicate) iterator nie zwróci już kolejnej wartości.

input = [1, 2, 4, 6, 2, 1]
result = takewhile(lambda x: x < 3, input)
print(list(result))  # [1, 2]

10. filterfalse

Tworzy iterator, który zwraca tylko wartości niespełniające warunku z parametru predicate. Jeśli argument predicate ma wartość None, wtedy zwracane są tylko wartości, które tłumaczy się jako boolowsko fałszywe.

Przykładowe odfiltrowanie ujemnych cyfr:

input = [1, -1, -2, 4, 5, -6]
result = filterfalse(lambda x: x < 0, input)
print(list(result))  # [1, 4, 5]

11. groupby

Jak sama nazwa wskazuje, funkcja ta grupuje kolejne elementy z przekazanego zbioru. Zwrócone przez nią grupy są iteratorem, więc jeśli chcemy je użyć więcej niż tylko raz, warto przekształcić je w listę.

input = 'AAAABBBCCDYYY'
for key, group in groupby(input):
    print(f'{key} - {list(group)}')
# A - ['A', 'A', 'A', 'A']
# B - ['B', 'B', 'B']
# C - ['C', 'C']
# D - ['D']
# Y - ['Y', 'Y', 'Y']

Bardzo ważne, aby dane wejściowe były posortowane. Dlaczego? Ponieważ groupby tworzy nową grupę jeśli kolejna wartość konsumowana przez niego jest inna niż poprzednia. Przekazując nieposortowane dane, wynik będzie jak poniżej:

input = 'ABBABBAA'
for key, group in groupby(input):
    print(f'{key} - {list(group)}')
# A - ['A']
# B - ['B', 'B']
# A - ['A']
# B - ['B', 'B']
# A - ['A', 'A']

Czyli otrzymaliśmy zdublowane grupy. W odróżnieniu GROUP BY znany z SQL potrafi pogrupować dane nawet jeśli dane nie są posortowane, podczas gdy groupby potrzebuje posortowanych danych na wejściu.

groupby akceptuje argument key do określenia wartości, której użyć do grupowania danych.

input = [
    {'name': 'Bob', 'level': 1},
    {'name': 'Alfred', 'level': 3},
    {'name': 'Greg', 'level': 4},
    {'name': 'Leon', 'level': 4},
]

for key, group in groupby(input, key=lambda x: x['level']):
    print(f'{key} - {list(group)}')
# 1 - [{'name': 'Bob', 'level': 1}]
# 3 - [{'name': 'Alfred', 'level': 3}]
# 4 - [{'name': 'Greg', 'level': 4}, {'name': 'Leon', 'level': 4}

12. islice

Zachowuje się analogicznie do dzielenia np. list (islice zwraca iterator). Przekazujemy iterable wraz z opcjonalnymi: indeksem początkowym, końcowym oraz krokiem.

Warto pamiętać, że jeśli podamy tylko jeden argument oprócz samego iterable, to będzie on oznaczał pozycję końcową zwróconego iteratora.

result = islice(range(5), 3)
print(list(result))  # [0, 1, 2]

Dodając kolejny argument, określamy kolejno pozycję początkową i końcową (w tej kolejności).

result = islice(range(5), 2, 4)
print(list(result))  # [2, 3]

Dodatkowo możemy określić jeszcze krok/skok oznaczający co która wartość ma być zwrócona.

result = islice(range(10), 1, 9, 3)
print(list(result))  # [1, 4, 7]

13. starmap

Działa bardzo podobnie do znanej funkcji map z jedną delikatną różnicą, że jako drugi przyjmuje iterowalny zbiór iterable 🙂

input = [(1, 0, -1), (4, 2, 3), (5, 2, 7)]
result = starmap(max, input)
print(list(result)) # [1, 4, 7]

Czyli zamiast wykonywać za każdym razem funkcję map na kolejnej krotce (ang. tuple), przekazujemy zbiór krotek. starmap iteruje po każdej krotce i wykonuje na nich – w naszym przypadku – funkcję wybierającą największą wartość.

14. tee

Jak wiemy z wpisu o iteratorach, przez stworzony iterator można przejść tylko raz. Po wyczerpaniu należy stworzyć kolejny jeśli chcemy skorzystać z niego znowu. Funkcja tee pozwala stworzyć za jednym zamachem X niezależnych iteratorów z pojedyńczego iterable. Domyślnie tworzy dwa iteratory.

text = 'abc'
iterator1, iterator2 = tee(text)
print(next(iterator1)) # a
print(next(iterator1)) # b
print(next(iterator1)) # c
print(next(iterator2)) # a
print(next(iterator2)) # b
print(next(iterator2)) # c

Ważne: dokumentacja mówi, że po użyciu funkcji tee obiekt iterable nie powinien być użyty nigdzie indziej. Może to prowadzić do niespodziewanego i błędnego zachowania.

Dodatkowo oficjalna dokumentacja zawiera informację, ostrzegającą o tym, że tee() nie jest bezpieczna w użyciu przy wątkach. Równoczesne używanie iteratorów zwróconych przez tee może prowadzić do zgłoszenia wyjątku RuntimeError.

15. zip_longest

Jeśli używałeś/aś zip do łączenia dwóch iterable razem to być może wiesz, że ilość elementów takiego iteratora równa się najkrótszemu z iterable. Dowód:

input_1 = [1, 2, 3, 4]
input_2 = 'abc'
result = zip(input_1, input_2)
print(list(result))  # [(1, 'a'), (2, 'b'), (3, 'c')]

Jak widzisz, iterator zwrócił tylko trzy elementy, bo właśnie tyle ma najkrótszy zbiór (input_2). zip_longest natomiast zachowa się odwrotnie. Zwróci ilość elementów równą najdłuższemu z przekazanych iterable. Domyślnie w miejsce brakujących wartości z krótszego zbioru wstawi wartość None. Możemy przekazać inną dowolną wartość, dodając argument fillvalue.

input_1 = [1, 2, 3, 4]
input_2 = 'abc'
result = zip_longest(input_1, input_2, fillvalue=0)
print(list(result))  # [(1, 'a'), (2, 'b'), (3, 'c'), (4, 0)]

Iteratory kombinatoryczne

16. product

Służy do obliczenia iloczynu kartezjańskiego. Jeśli iterable wejściowe są posortowane, to na wyjściu elementy również będą posortowane.

Przedstawię poniżej bardzo fajny przykład zaczerpnięty stąd. W dość łatwy sposób generujemy wszystkie karty do gry 🙂

FACE_CARDS = ("J", "Q", "K", "A")
SUITS = ("♥", "♦", "♣", "♠")

DECK = list(
    product(
        chain(range(2, 11), FACE_CARDS),
        SUITS,
    )
)

for card in DECK:
    print("{:>2}{}".format(*card), end=" ")
    if card[1] == SUITS[-1]:  # przeskocz do nowej linii jeśli kolor to pik
        print()

#  2♥  2♦  2♣  2♠
#  3♥  3♦  3♣  3♠
#  4♥  4♦  4♣  4♠
#  5♥  5♦  5♣  5♠
#  6♥  6♦  6♣  6♠
#  7♥  7♦  7♣  7♠
#  8♥  8♦  8♣  8♠
#  9♥  9♦  9♣  9♠
# 10♥ 10♦ 10♣ 10♠
#  J♥  J♦  J♣  J♠
#  Q♥  Q♦  Q♣  Q♠
#  K♥  K♦  K♣  K♠
#  A♥  A♦  A♣  A♠

Dodatkowo w przykładzie powyżej użyliśmy opisany wcześniejchain, aby w prosty sposób dodać do siebie karty numerowane z waletem, damą, królem oraz asem. Dzięki product() otrzymaliśmy kombinacje wszystkich figur z kolorami.

Dodatkowo product pozwala na określenie ile razy zduplikować przekazane jako argument zbiory. Służy do tego parametr repeat. Oto prosty przykład obrazujący zachowanie tego parametru:

result = product(range(2))
print(list(result))  # [(0,), (1,)]

Jak widać, przekazaliśmy tylko jeden iterable, więc otrzymujemy właściwie te same wartości, które przekazaliśmy. Teraz dodajmy parametr repeat=2:

result = product(range(2), repeat=2)
print(list(result))  # [(0, 0), (0, 1), (1, 0), (1, 1)]

Od razu lepiej. Jeśli chcemy obliczyć iloczyn kartezjański, używając takich samych zbiorów, warto pamiętać o tym parametrze.

17. permutations

Dzięki tej funkcji otrzymamy permutacje przekazanego zbioru. Domyślnie długość każdej permutacji równa jest długości przekazanego iterable. Przekazując argument r możemy nim manipulować.

result = permutations('ABC')
print(list(result))  # [('A', 'B', 'C'), ('A', 'C', 'B'), ('B', 'A', 'C'), ('B', 'C', 'A'), ('C', 'A', 'B'), ('C', 'B', 'A')]

result = permutations('ABC', r=2)
print(list(result))  # [('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'C'), ('C', 'A'), ('C', 'B')]

18. combinations

Oblicza kombinacje dla podanego zbioru. W tym przypadku wymagane jest przekazanie argumentu r określającego długość kombinacji.

result = combinations('ABC', r=2)
print(list(result))  # [('A', 'B'), ('A', 'C'), ('B', 'C')]

19. combinations_with_replacement

Zachowuje się analogicznie jak opisane powyżej combinations z tą różnicą, że pozwala na powtórzenie tego samego elementu.

result = combinations_with_replacement('ABC', r=2)
print(list(result))  # [('A', 'A'), ('A', 'B'), ('A', 'C'), ('B', 'B'), ('B', 'C'), ('C', 'C')]

Nie tylko Python Itertools… 🙂

Jeśli nie znalazłeś jeszcze tego, czego szukasz w bibliotece standardowej Python Itertools, być może znajdziesz to w paczce more-itertools. Zawiera duuużo różnych i przede wszystkim wydajnych narzędzi, które mogą nam zaoszczędzić sporo czasu. Warto przejść przez listę funkcji, bo być może znajdziesz coś dla siebie. Projekt na czas pisania tego posta jest aktywnie rozwijany.

Artykuły zagraniczne fajnie opisujące bibliotekę itertools

  • https://realpython.com/python-itertools/
  • https://florian-dahlitz.de/blog/introduction-to-itertools

Bądź na bieżąco z blogiem!

Bądź na bieżąco z kolejnymi metariałami publikowanymi na tym blogu i innymi użytecznymi (moim zdaniem) materiałami zapisując się do newslettera 🙂

Jako prezent otrzymasz ode mnie mini e-book „Jak zacząć działać w Open Source i rozwijać swoje umiejętności” 🙂