Сообщить об ошибке.

Трюки и рецепты использования модуля itertools

В этом разделе показаны рецепты создания расширенного набора инструментов с использованием существующих функция модуля itertools в качестве строительных блоков.

  • Предложенные функции предлагают такую ​​же высокую производительность, как и функции в модулях.
  • Малое потребление памяти достигается обработкой элементов по одному, а не переносом всей итерируемой последовательности в память.
  • Высокая скорость сохраняется за счет предпочтения "векторизованных" блоков над использованием циклов for и генераторов, которые влекут за собой некоторые издержки.

Содержание:

Получить первые n элементов итератора/списка в виде списка:

import itertools

def take(n, iterable):
    # take(3, 'ABCDEFG') --> A B C
    return list(itertools.islice(iterable, n))

>>> l = [2, 3, 5, 4, 7, 6, 8, 9]
>>> take(3, l)
# [2, 3, 5]
>>> take(5, l)
# [2, 3, 5, 4, 7]

Получить n последних элементов итератора/списка:

import collections

def tail(n, iterable):
    # tail(3, 'ABCDEFG') --> E F G
    return iter(collections.deque(iterable, maxlen=n))

>>> l = [2, 3, 5, 4, 7, 6, 8, 9]
>>> z = tail(3, l)
>>> list(z)
# [6, 8, 9]
>>> z = tail(5, l)
>>> list(z)
# [4, 7, 6, 8, 9]

Получить n-й элемент или значение по умолчанию:

import itertools

def nth(iterable, n, default=None):
    return next(itertools.islice(iterable, n, None), default)

>>> l = [2, 3, 5, 7]
>>> l[8]
# Traceback (most recent call last):
#   File "<stdin>", line 1, in <module>
# IndexError: list index out of range
>>> nth(l, 8, 0)
# 0
>>> nth(l, 3, 0)
# 7

True, если все элементы списка равны друг другу:

import itertools

def all_equal(iterable):
    g = itertools.groupby(iterable)
    return next(g, True) and not next(g, False)

Посчитать истинные значения в списке/итераторе:

def quantify(iterable, pred=bool):
    return sum(map(pred, iterable))

>>> x = list(range(5))
>>> x = x*3
>>> x
# [0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4]
>>> quantify(x, bool)    # цифра 0 эквивалентна False
# 12

# сколько раз встречается цифра 2 в списке
>>> quantify(x, lambda z: z==2)
# 3

Лениво вернуть элементы списка n раз:

import itertools

def ncycles(iterable, n):
    return itertools.chain.from_iterable(itertools.repeat(tuple(iterable), n))

>>> l = [1, 3, 5, 7]
>>> ncycles(l, 1)
# <itertools.chain object at 0x7fc4f9e40e80>
>>> z = ncycles(l, 3)
>>> list(z)
# [1, 3, 5, 7, 1, 3, 5, 7, 1, 3, 5, 7]

Сумма произведений элементов 2-х списков:

import operator

def dotproduct(vec1, vec2):
    return sum(map(operator.mul, vec1, vec2))

>>> l1 = [2, 3, 5, 5]
>>> l2 = [4, 6, 8, 9]
>>> dotproduct(l1, l2)
# 111

>>> l1 = [2, 3]
>>> l2 = [4, 6, 8, 9]
>>> dotproduct(l1, l2)
26

Раскрыть один уровень вложенного списка:

import itertools

def flatten(list_of_lists):
    return itertools.chain.from_iterable(list_of_lists)

>>> l = [[2, 3, 5], [4, 7, 6, 8], [9]]
>>> z = flatten(l)
>>> list(z)
# [2, 3, 5, 4, 7, 6, 8, 9]

Повторить вызов функции с указанными аргументами:

import itertools

def repeatfunc(func, times=None, *args):
    if times is None:
        return itertools.starmap(func, itertools.repeat(args))
    return itertools.starmap(func, itertools.repeat(args, times))

>>> import random
# список случайных значений
>>> z = repeatfunc(random.randint, 15, 1, 50)
>>> list(z)
# [40, 31, 38, 41, 11, 20, 32, 20, 20, 9, 42, 38, 1, 48, 9]

Объединить элементы списка => (s0,s1), (s1,s2), (s2, s3), ...:

import itertools

def pairwise(iterable):
    a, b = itertools.tee(iterable)
    next(b, None)
    return zip(a, b)

>>> l = [2, 3, 5, 4, 7, 6, 8, 9]
>>> z = pairwise(l)
>>> list(z)
# [(2, 3), (3, 5), (5, 4), (4, 7), (7, 6), (6, 8), (8, 9)]

Объединить элементы списка в кортежи по n элементов:

import itertools

def grouper(iterable, n, fillvalue=None):
    # grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx"
    args = [iter(iterable)] * n
    return itertools.zip_longest(*args, fillvalue=fillvalue)

>>> l = [4, 7, 6, 2, 10, 5, 2, 2, 7, 3]
>>> a = grouper(l, 3)
>>> list(a)
# [(4, 7, 6), (2, 10, 5), (2, 2, 7), (3, None, None)]

# делим список пополам
>>> n = len(l)
>>> a = grouper(l, n//2)
>>> list(a)
# [(4, 7, 6, 2, 10), (5, 2, 2, 7, 3)]

Round Robin:

import itertools

def roundrobin(*iterables):
    "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
    num_active = len(iterables)
    nexts = itertools.cycle(iter(it).__next__ for it in iterables)
    while num_active:
        try:
            for next in nexts:
                yield next()
        except StopIteration:
            # Remove the iterator we just exhausted from the cycle.
            num_active -= 1
            nexts = itertools.cycle(itertools.islice(nexts, num_active))

>>> z = roundrobin('abc','d', 'ef')
>>> list(z)
['a', 'd', 'e', 'b', 'f', 'c']

Разделить элементы списка на ложные и истинные:

import itertools

def partition(pred, iterable):
    # partition(is_odd, range(10)) --> 0 2 4 6 8   and  1 3 5 7 9
    t1, t2 = itertools.tee(iterable)
    return itertools.filterfalse(pred, t1), filter(pred, t2)

>>> a = partition(lambda x: x%2, range(10))
>>> one, two = list(a[0]), list(a[1])
>>> one
[0, 2, 4, 6, 8]
>>> two
[1, 3, 5, 7, 9]

Список уникальных элементов списка, сохраняющих порядок:

import itertools

def unique_everseen(iterable, key=None):
    # unique_everseen('AAAABBBCCDAABBB') --> A B C D
    # unique_everseen('ABBCcAD', str.lower) --> A B C D
    seen = set()
    seen_add = seen.add
    if key is None:
        for element in itertools.filterfalse(seen.__contains__, iterable):
            seen_add(element)
            yield element
    else:
        for element in iterable:
            k = key(element)
            if k not in seen:
                seen_add(k)
                yield element

>>> z = unique_everseen('AAAABBBCCDAABBB')
>>> list(z)
# ['A', 'B', 'C', 'D']
>>> z = unique_everseen('AAAABbBCcDAAaBBbB')
>>> list(z)
# ['A', 'B', 'b', 'C', 'c', 'D', 'a']
>>> z = unique_everseen('AAAABbBCcDAAaBBbB', str.lower)
>>> list(z)
# ['A', 'B', 'C', 'D']

Удаление повторов соседних элементов списка:

import itertools, operator

def unique_justseen(iterable, key=None):
    # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
    # unique_justseen('ABBCcAD', str.lower) --> A B C A D
    return map(next, map(operator.itemgetter(1), itertools.groupby(iterable, key)))

>>> z = unique_justseen('AAAABBBCCDAABBB')
>>> list(z)
# ['A', 'B', 'C', 'D', 'A', 'B']
>>> z = unique_justseen('AAAABBBCCDAAaBBbB')
>>> list(z)
# ['A', 'B', 'C', 'D', 'A', 'a', 'B', 'b', 'B']
>>> z = unique_justseen('AAAABBBCCDAAaBBbB', str.lower)
>>> list(z)
# ['A', 'B', 'C', 'D', 'A', 'B']

Повторно вызывает функцию, пока не возникнет исключение:

import itertools

def iter_except(func, exception, first=None):
    """ Преобразует интерфейс вызова до exception в интерфейс итератора.
    Examples:
        # приоритетный итератор очереди
        iter_except(functools.partial(heappop, h), IndexError)
        # неблокирующий итератор dict
        iter_except(d.popitem, KeyError)
        # Неблокирующий итератор deque
        iter_except(d.popleft, IndexError)
        # цикл по очереди
        iter_except(q.get_nowait, Queue.Empty)
        # Неблокирующий итератор множества
        iter_except(s.pop, KeyError)
    """
    try:
        if first is not None:
            yield first()
        while True:
            yield func()
    except exception:
        pass

Получить первое истинное значение из списка:

Если истинное значение не найдено, возвращает default.Если pred не равно None, возвращает первый элемент для которого pred(item) верно.

import itertools

def first_true(iterable, default=False, pred=None):
    # first_true([a,b,c], x) --> a or b or c or x
    # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
    return next(filter(pred, iterable), default)

>>> first_true([0,0,0], 121)
# 121
>>> first_true([0,23,0, 50])
# 23
>>> first_true([0,23,0], 121, lambda x: x==0)
# 0
>>> first_true([0,23,0], 121, lambda x: x==5)
# 121

Случайно выбрать по n элементов из нескольких списков:

import random

def random_product(*args, repeat=1):
    pools = [tuple(pool) for pool in args] * repeat
    return tuple(random.choice(pool) for pool in pools)

>>> l1 = [2, 3, 5, 5]
>>> l2 = [4, 6, 8, 9]
>>> l3 = [10, 5, 6, 3]
>>> random_product(l1, l2, l3) 
# (5, 8, 10)
>>> random_product(l1, l2, l3, repeat=3) 
# (2, 6, 6, 5, 8, 6, 3, 4, 10)

Перемешать список случайным образом:

import random

def random_permutation(iterable, r=None):
    pool = tuple(iterable)
    r = len(pool) if r is None else r
    return tuple(random.sample(pool, r))

>>> l = [2, 3, 5, 4, 6, 8, 9]
>>> random_permutation(l)
# (6, 2, 8, 9, 4, 5, 3)

Выбрать n элементов из списка случайным образом:

import random

def random_combination(iterable, r):
    pool = tuple(iterable)
    n = len(pool)
    indices = sorted(random.sample(range(n), r))
    return tuple(pool[i] for i in indices)

>>> l = [2, 3, 5, 4, 7, 6, 8, 9]
>>> random_combination(l, 3)
# (5, 8, 9)
>>> random_combination(l, 5)
# (3, 5, 7, 6, 9)