itertools — Функции создания итераторов для эффективного циклирования
Этот модуль реализует ряд итераторов, вдохновлённых конструкциями из APL, Haskell и SML. Каждый из них переработан в форме, подходящей для Python.
Модуль стандартизирует набор быстрых и экономичных инструментов, которые полезны сами по себе или в комбинации. Вместе они образуют «алгебру итераторов», что позволяет создавать специализированные инструменты лаконично и эффективно в чистом Python.
Например, SML предоставляет инструмент табулирования: tabulate(f) , который генерирует последовательность f(0), f(1), .... То же самое можно достичь в Python, объединив map() и count(), чтобы получить map(f, count()).
Эти инструменты и их встроенные аналоги также хорошо работают с высокоскоростными функциями в модуле operator. Например, оператор умножения можно применить к двум векторам для формирования эффективного скалярного произведения: sum(starmap(operator.mul, zip(vec1, vec2, strict=True))).
Бесконечные итераторы:
Итератор | Аргументы | Результаты | Пример |
|---|---|---|---|
start, [step] | start, start+step, start+2*step, … |
| |
p | p0, p1, … plast, p0, p1, … |
| |
elem [,n] | elem, elem, elem, … бесконечно или до n раз |
|
Итераторы, завершающиеся на самой короткой входной последовательности:
Итератор | Аргументы | Результаты | Пример |
|---|---|---|---|
p [,func] | p0, p0+p1, p0+p1+p2, … |
| |
p, q, … | p0, p1, … plast, q0, q1, … |
| |
iterable | p0, p1, … plast, q0, q1, … |
| |
data, selectors | (d[0] if s[0]), (d[1] if s[1]), … |
| |
pred, seq | seq[n], seq[n+1], начиная с момента, когда pred возвращает False |
| |
pred, seq | элементы из seq, где pred(elem) ложно |
| |
iterable[, key] | под-итераторы, сгруппированные по значению key(v) | ||
seq, [start,] stop [, step] | элементы из seq[start:stop:step] |
| |
iterable | (p[0], p[1]), (p[1], p[2]) |
| |
func, seq | func(*seq[0]), func(*seq[1]), … |
| |
pred, seq | seq[0], seq[1], до тех пор, пока pred не вернёт False |
| |
it, n | it1, it2, … itn разделяет один итератор на n | ||
p, q, … | (p[0], q[0]), (p[1], q[1]), … |
|
Комбинаторные итераторы:
Итератор | Аргументы | Результаты |
|---|---|---|
p, q, … [repeat=1] | декартово произведение, эквивалентно вложенному циклу for | |
p[, r] | кортежи длины r, все возможные упорядочения, без повторяющихся элементов | |
p, r | кортежи длины r, в отсортированном порядке, без повторяющихся элементов | |
p, r | кортежи длины r, в отсортированном порядке, с повторяющимися элементами |
Примеры | Результаты |
|---|---|
|
|
|
|
|
|
|
|
Функции itertools
Следующие функции модуля строят и возвращают итераторы. Некоторые предоставляют потоки бесконечной длины, поэтому к ним следует обращаться только в функциях или циклах, которые обрезают поток.
-
itertools.accumulate(iterable[, func, *, initial=None]) -
Создайте итератор, возвращающий накопленные суммы или накопленные результаты других бинарных функций (указанных с помощью необязательного аргумента func).
Если func указан, он должен быть функцией двух аргументов. Элементы входного iterable могут быть любого типа, который может быть принят в качестве аргументов для func. (Например, с помощью стандартной операции сложения, элементы могут быть любого типа, поддерживающего сложение, включая
DecimalилиFraction.)Обычно количество выводимых элементов соответствует входному iterable. Однако, если указан ключевой аргумент initial, накопление начинается со значения initial, так что выходной поток имеет на один элемент больше, чем входной iterable.
Приблизительно эквивалентно:
def accumulate(iterable, func=operator.add, *, initial=None): 'Return running totals' # accumulate([1,2,3,4,5]) --> 1 3 6 10 15 # accumulate([1,2,3,4,5], initial=100) --> 100 101 103 106 110 115 # accumulate([1,2,3,4,5], operator.mul) --> 1 2 6 24 120 it = iter(iterable) total = initial if initial is None: try: total = next(it) except StopIteration: return yield total for element in it: total = func(total, element) yield totalАргумент func может быть использован для различных целей. Он может быть установлен в
min()для вычисления текущего минимума,max()для вычисления текущего максимума илиoperator.mul()для вычисления текущего произведения. Таблицы амортизации могут быть построены путем накопления процентов и применения платежей:>>> data = [3, 4, 6, 2, 1, 9, 0, 7, 5, 8] >>> list(accumulate(data, operator.mul)) # running product [3, 12, 72, 144, 144, 1296, 0, 0, 0, 0] >>> list(accumulate(data, max)) # running maximum [3, 4, 6, 6, 6, 9, 9, 9, 9, 9] # Amortize a 5% loan of 1000 with 4 annual payments of 90 >>> cashflows = [1000, -90, -90, -90, -90] >>> list(accumulate(cashflows, lambda bal, pmt: bal*1.05 + pmt)) [1000, 960.0, 918.0, 873.9000000000001, 827.5950000000001]
См.
functools.reduce()для аналогичной функции, которая возвращает только конечное накопленное значение.Новое в версии 3.2.
Изменено в версии 3.3: Добавлен необязательный параметр func.
Изменено в версии 3.8: Добавлен необязательный параметр initial.
-
itertools.chain(*iterables) -
Создайте итератор, возвращающий элементы из первого итерируемого объекта, пока он не будет исчерпан, затем переходит к следующему итерируемому объекту, пока все итерируемые объекты не будут исчерпаны. Используется для обработки последовательных последовательностей как одной последовательности. Приблизительно эквивалентно:
def chain(*iterables): # chain('ABC', 'DEF') --> A B C D E F for it in iterables: for element in it: yield element
-
classmethod chain.from_iterable(iterable) -
Альтернативный конструктор для
chain(). Получает входные данные из единственного итерируемого аргумента, который вычисляется лениво. Приблизительно эквивалентно:def from_iterable(iterables): # chain.from_iterable(['ABC', 'DEF']) --> A B C D E F for it in iterables: for element in it: yield element
-
itertools.combinations(iterable, r) -
Возвращает подпоследовательности длины r элементов из входного iterable.
Кортежи комбинаций генерируются в лексикографическом порядке, согласно порядку входного iterable. Таким образом, если входной iterable отсортирован, выходные кортежи будут генерироваться в отсортированном порядке.
Элементы обрабатываются как уникальные, основываясь на их позиции, а не на их значении. Таким образом, если входные элементы уникальны, в каждой комбинации не будет повторяющихся значений.
Приблизительно эквивалентно:
def combinations(iterable, r): # combinations('ABCD', 2) --> AB AC AD BC BD CD # combinations(range(4), 3) --> 012 013 023 123 pool = tuple(iterable) n = len(pool) if r > n: return indices = list(range(r)) yield tuple(pool[i] for i in indices) while True: for i in reversed(range(r)): if indices[i] != i + n - r: break else: return indices[i] += 1 for j in range(i+1, r): indices[j] = indices[j-1] + 1 yield tuple(pool[i] for i in indices)Код для
combinations()также может быть выражен как подпоследовательностьpermutations()после фильтрации записей, где элементы не отсортированы (согласно их позиции в входном наборе):def combinations(iterable, r): pool = tuple(iterable) n = len(pool) for indices in permutations(range(n), r): if sorted(indices) == list(indices): yield tuple(pool[i] for i in indices)Количество возвращаемых элементов равно
n! / r! / (n-r)!при0 <= r <= nили нулю приr > n.
-
itertools.combinations_with_replacement(iterable, r) -
Возвращает подпоследовательности длины r элементов из входного iterable, позволяя повторять отдельные элементы более одного раза.
Кортежи комбинаций генерируются в лексикографическом порядке, согласно порядку входного iterable. Таким образом, если входной iterable отсортирован, выходные кортежи будут генерироваться в отсортированном порядке.
Элементы обрабатываются как уникальные, основываясь на их позиции, а не на их значении. Таким образом, если входные элементы уникальны, сгенерированные комбинации также будут уникальными.
Приблизительно эквивалентно:
def combinations_with_replacement(iterable, r): # combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC pool = tuple(iterable) n = len(pool) if not n and r: return indices = [0] * r yield tuple(pool[i] for i in indices) while True: for i in reversed(range(r)): if indices[i] != n - 1: break else: return indices[i:] = [indices[i] + 1] * (r - i) yield tuple(pool[i] for i in indices)Код для
combinations_with_replacement()также может быть выражен как подпоследовательностьproduct()после фильтрации записей, где элементы не отсортированы (согласно их позиции в входном наборе):def combinations_with_replacement(iterable, r): pool = tuple(iterable) n = len(pool) for indices in product(range(n), repeat=r): if sorted(indices) == list(indices): yield tuple(pool[i] for i in indices)Количество возвращаемых элементов равно
(n+r-1)! / r! / (n-1)!приn > 0.Новое в версии 3.1.
-
itertools.compress(data, selectors) -
Создайте итератор, который фильтрует элементы из data, возвращая только те, у которых соответствующий элемент в selectors вычисляется как
True. Остановка происходит, когда либо итерируемые объекты data или selectors были исчерпаны. Приблизительно эквивалентно:def compress(data, selectors): # compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F return (d for d, s in zip(data, selectors) if s)Новое в версии 3.1.
-
itertools.count(start=0, step=1) -
Создайте итератор, возвращающий равномерно распределённые значения, начиная с числа start. Часто используется в качестве аргумента для
map()для генерации последовательных точек данных. Также используется сzip()для добавления номеров последовательности. Приблизительно эквивалентно:def count(start=0, step=1): # count(10) --> 10 11 12 13 14 ... # count(2.5, 0.5) --> 2.5 3.0 3.5 ... n = start while True: yield n n += stepПри подсчёте с плавающей точкой, для повышения точности иногда можно использовать код с умножением, например:
(start + step * i for i in count()).Изменено в версии 3.1: Добавлен аргумент step и разрешены нецелочисленные аргументы.
-
itertools.cycle(iterable) -
Создайте итератор, возвращающий элементы из итерируемого объекта и сохраняющий копию каждого. После исчерпания итерируемого объекта, возвращает элементы из сохранённой копии. Повторяется бесконечно. Приблизительно эквивалентно:
def cycle(iterable): # cycle('ABCD') --> A B C D A B C D A B C D ... saved = [] for element in iterable: yield element saved.append(element) while saved: for element in saved: yield elementОбратите внимание, что этот элемент инструментария может потребовать значительного дополнительного хранения (в зависимости от длины итерируемого объекта).
-
itertools.dropwhile(predicate, iterable) -
Создайте итератор, пропускающий элементы из итерируемого объекта, пока предикат истинен; после этого возвращает каждый элемент. Обратите внимание, что итератор не генерирует никакого вывода, пока предикат не станет ложным, поэтому у него может быть длительный начальный этап. Приблизительно эквивалентно:
def dropwhile(predicate, iterable): # dropwhile(lambda x: x<5, [1,4,6,4,1]) --> 6 4 1 iterable = iter(iterable) for x in iterable: if not predicate(x): yield x break for x in iterable: yield x
-
itertools.filterfalse(predicate, iterable) -
Создайте итератор, фильтрующий элементы из итерируемого объекта, возвращая только те, для которых предикат ложен. Если predicate
None, возвращает элементы, которые ложны. Приблизительно эквивалентно:def filterfalse(predicate, iterable): # filterfalse(lambda x: x%2, range(10)) --> 0 2 4 6 8 if predicate is None: predicate = bool for x in iterable: if not predicate(x): yield x
-
itertools.groupby(iterable, key=None) -
Создайте итератор, возвращающий последовательные ключи и группы из iterable. key — функция, вычисляющая значение ключа для каждого элемента. Если не указана или равна
None, key по умолчанию — функция идентичности и возвращает элемент без изменений. Обычно, итерируемый объект должен быть предварительно отсортирован по той же функции ключа.Действие
groupby()похоже на командуuniqфильтрации в Unix. Оно генерирует разрыв или новую группу каждый раз, когда значение функции ключа меняется (поэтому обычно необходимо отсортировать данные с использованием той же функции ключа). Это отличается от SQL’s GROUP BY, которая агрегирует общие элементы независимо от порядка ввода.Возвращаемая группа сама является итератором, который разделяет базовый итерируемый объект с
groupby(). Поскольку источник общий, когда объектgroupby()перемещается вперёд, предыдущая группа больше не видна. Поэтому, если эти данные нужны позже, их следует сохранить в списке:groups = [] uniquekeys = [] data = sorted(data, key=keyfunc) for k, g in groupby(data, keyfunc): groups.append(list(g)) # Store group iterator as a list uniquekeys.append(k)groupby()примерно эквивалентно:class groupby: # [k for k, g in groupby('AAAABBBCCDAABBB')] --> A B C D A B # [list(g) for k, g in groupby('AAAABBBCCD')] --> AAAA BBB CC D def __init__(self, iterable, key=None): if key is None: key = lambda x: x self.keyfunc = key self.it = iter(iterable) self.tgtkey = self.currkey = self.currvalue = object() def __iter__(self): return self def __next__(self): self.id = object() while self.currkey == self.tgtkey: self.currvalue = next(self.it) # Exit on StopIteration self.currkey = self.keyfunc(self.currvalue) self.tgtkey = self.currkey return (self.currkey, self._grouper(self.tgtkey, self.id)) def _grouper(self, tgtkey, id): while self.id is id and self.currkey == tgtkey: yield self.currvalue try: self.currvalue = next(self.it) except StopIteration: return self.currkey = self.keyfunc(self.currvalue)
-
itertools.islice(iterable, stop) - itertools.islice(iterable, start, stop[, step])
-
Создайте итератор, возвращающий выбранные элементы из итерируемого объекта. Если start не равно нулю, то элементы из итерируемого объекта пропускаются до достижения start. После этого элементы возвращаются последовательно, если step не установлен выше единицы, что приводит к пропуску элементов. Если stop
None, итерация продолжается до тех пор, пока итератор не будет исчерпан, если вообще. В противном случае она останавливается в указанной позиции.Если start
None, итерация начинается с нуля. Если stepNone, шаг по умолчанию равен единице.В отличие от обычного среза,
islice()не поддерживает отрицательные значения для start, stop или step. Может использоваться для извлечения связанных полей из данных, где внутренняя структура была уплощена (например, многострочный отчёт может перечислять поле имени на каждой третьей строке).Приблизительно эквивалентно:
def islice(iterable, *args): # islice('ABCDEFG', 2) --> A B # islice('ABCDEFG', 2, 4) --> C D # islice('ABCDEFG', 2, None) --> C D E F G # islice('ABCDEFG', 0, None, 2) --> A C E G s = slice(*args) start, stop, step = s.start or 0, s.stop or sys.maxsize, s.step or 1 it = iter(range(start, stop, step)) try: nexti = next(it) except StopIteration: # Consume *iterable* up to the *start* position. for i, element in zip(range(start), iterable): pass return try: for i, element in enumerate(iterable): if i == nexti: yield element nexti = next(it) except StopIteration: # Consume to *stop*. for i, element in zip(range(i + 1, stop), iterable): pass
-
itertools.pairwise(iterable) -
Возвращает последовательные перекрывающиеся пары, взятые из входного итерируемого объекта.
Количество 2-кортежей в итераторе вывода будет на один меньше, чем количество входных значений. Он будет пустым, если входной итерируемый объект содержит меньше двух значений.
Приблизительно эквивалентно:
def pairwise(iterable): # pairwise('ABCDEFG') --> AB BC CD DE EF FG a, b = tee(iterable) next(b, None) return zip(a, b)Новое в версии 3.10.
-
itertools.permutations(iterable, r=None) -
Возвращает последовательные перестановки длины r элементов в итерируемом объекте.
Если r не указано или равно
None, то r по умолчанию равно длине итерируемого объекта, и генерируются все возможные перестановки полной длины.Кортежи перестановок генерируются в лексикографическом порядке в соответствии с порядком входного итерируемого объекта. Таким образом, если входной итерируемый объект отсортирован, выходные кортежи будут генерироваться в отсортированном порядке.
Элементы обрабатываются как уникальные на основе их позиции, а не значения. Таким образом, если входные элементы уникальны, в перестановке не будет повторяющихся значений.
Приблизительно эквивалентно:
def permutations(iterable, r=None): # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC # permutations(range(3)) --> 012 021 102 120 201 210 pool = tuple(iterable) n = len(pool) r = n if r is None else r if r > n: return indices = list(range(n)) cycles = list(range(n, n-r, -1)) yield tuple(pool[i] for i in indices[:r]) while n: for i in reversed(range(r)): cycles[i] -= 1 if cycles[i] == 0: indices[i:] = indices[i+1:] + indices[i:i+1] cycles[i] = n - i else: j = cycles[i] indices[i], indices[-j] = indices[-j], indices[i] yield tuple(pool[i] for i in indices[:r]) break else: returnКод для
permutations()также может быть выражен как подпоследовательностьproduct(), отфильтрованная для исключения записей с повторяющимися элементами (те, которые из одной позиции в входном наборе):def permutations(iterable, r=None): pool = tuple(iterable) n = len(pool) r = n if r is None else r for indices in product(range(n), repeat=r): if len(set(indices)) == r: yield tuple(pool[i] for i in indices)Количество возвращаемых элементов равно
n! / (n-r)!когда0 <= r <= nили нулю, когдаr > n.
-
itertools.product(*iterables, repeat=1) -
Декартово произведение входных итерируемых объектов.
Приблизительно эквивалентно вложенным циклам for в выражении генератора. Например,
product(A, B)возвращает то же самое, что и((x,y) for x in A for y in B).Вложенные циклы циклически изменяются, как одометр, при каждом проходе правого элемента. Этот шаблон создаёт лексикографический порядок, так что если входные итерируемые объекты отсортированы, кортежи произведения генерируются в отсортированном порядке.
Чтобы вычислить произведение итерируемого объекта с самим собой, укажите количество повторений с помощью необязательного ключевого аргумента repeat. Например,
product(A, repeat=4)означает то же самое, что иproduct(A, A, A, A).Эта функция примерно эквивалентна следующему коду, за исключением того, что фактическая реализация не строит промежуточные результаты в памяти:
def product(*args, repeat=1): # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111 pools = [tuple(pool) for pool in args] * repeat result = [[]] for pool in pools: result = [x+[y] for x in result for y in pool] for prod in result: yield tuple(prod)Перед запуском
product(), он полностью потребляет входные итерируемые объекты, сохраняя наборы значений в памяти для генерации произведений. Соответственно, он полезен только с конечными входами.
-
itertools.repeat(object[, times]) -
Создаёт итератор, который возвращает объект снова и снова. Бесконечно работает, если не указан аргумент times.
Приблизительно эквивалентно:
def repeat(object, times=None): # repeat(10, 3) --> 10 10 10 if times is None: while True: yield object else: for i in range(times): yield objectРаспространённое использование repeat — это подача потока постоянных значений в map или zip:
>>> list(map(pow, range(10), repeat(2))) [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
-
itertools.starmap(function, iterable) -
Создаёт итератор, вычисляющий функцию с аргументами, полученными из итерируемого объекта. Используется вместо
map(), когда параметры аргументов уже сгруппированы в кортежи из одного итерируемого объекта (когда данные были «предварительно сжаты»).Разница между
map()иstarmap()аналогична различию междуfunction(a,b)иfunction(*c). Приблизительно эквивалентно:def starmap(function, iterable): # starmap(pow, [(2,5), (3,2), (10,3)]) --> 32 9 1000 for args in iterable: yield function(*args)
-
itertools.takewhile(predicate, iterable) -
Создаёт итератор, возвращающий элементы из итерируемого объекта, пока предикат истинен. Приблизительно эквивалентно:
def takewhile(predicate, iterable): # takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4 for x in iterable: if predicate(x): yield x else: break
-
itertools.tee(iterable, n=2) -
Возвращает n независимых итераторов из одного итерируемого объекта.
Следующий код Python поясняет, что делает tee (хотя фактическая реализация более сложная и использует только один базовый FIFO-очередь):
def tee(iterable, n=2): it = iter(iterable) deques = [collections.deque() for i in range(n)] def gen(mydeque): while True: if not mydeque: # when the local deque is empty try: newval = next(it) # fetch a new value and except StopIteration: return for d in deques: # load it to all the deques d.append(newval) yield mydeque.popleft() return tuple(gen(d) for d in deques)После создания
tee(), исходный итерируемый объект не должен использоваться где-либо ещё; в противном случае итерируемый объект может быть продвинут, не сообщая об этом объектам tee.Итераторы
teeне потокобезопасны. Может быть поднято исключениеRuntimeErrorпри одновременном использовании итераторов, возвращённых одним вызовомtee(), даже если исходный итерируемый объект потокобезопасен.Этот итерируемый объект может потребовать значительного вспомогательного хранилища (в зависимости от того, сколько временных данных необходимо хранить). В целом, если один итератор использует большинство или все данные перед тем, как другой итератор начнётся, быстрее использовать
list()вместоtee().
-
itertools.zip_longest(*iterables, fillvalue=None) -
Создаёт итератор, агрегирующий элементы из каждого из итерируемых объектов. Если итерируемые объекты имеют неравную длину, пропущенные значения заполняются fillvalue. Итерация продолжается до тех пор, пока самый длинный итерируемый объект не будет исчерпан. Приблизительно эквивалентно:
def zip_longest(*args, fillvalue=None): # zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D- iterators = [iter(it) for it in args] num_active = len(iterators) if not num_active: return while True: values = [] for i, it in enumerate(iterators): try: value = next(it) except StopIteration: num_active -= 1 if not num_active: return iterators[i] = repeat(fillvalue) value = fillvalue values.append(value) yield tuple(values)Если один из итерируемых объектов потенциально бесконечен, то функция
zip_longest()должна быть обернута чем-то, что ограничивает число вызовов (например,islice()илиtakewhile()). Если не указано, fillvalue по умолчанию равноNone.
Рецепты itertools
Этот раздел показывает рецепты для создания расширенного набора инструментов с использованием существующих itertools в качестве строительных блоков.
Основная цель рецептов itertools — образовательная. Рецепты показывают различные способы мышления об отдельных инструментах — например, что chain.from_iterable относится к понятию сплющивания. Рецепты также дают идеи о способах объединения инструментов — например, о том, как compress() и range() могут работать вместе. Рецепты также показывают шаблоны использования itertools с модулями operator и collections, а также с встроенными itertools, такими как map(), filter(), reversed(), и enumerate().
Вторичная цель рецептов — служить инкубатором. accumulate(), compress(), и pairwise() itertools изначально были рецептами. В настоящее время рецепт iter_index() тестируется, чтобы определить, оправдает ли он себя.
Почти все эти рецепты и многие другие можно установить из проекта more-itertools на Python Package Index:
python -m pip install more-itertools
Многие рецепты предлагают такую же высокую производительность, как и лежащий в их основе набор инструментов. Превосходная производительность памяти сохраняется за счёт обработки элементов по одному, а не загрузки всего итерируемого объекта в память сразу. Объем кода поддерживается небольшим за счёт объединения инструментов в функциональном стиле, что помогает устранить временные переменные. Высокая скорость сохраняется за счёт предпочтения «векторизованных» строительных блоков перед использованием циклов for и генераторов, которые влекут за собой накладные расходы интерпретатора.
import collections
import math
import operator
import random
def take(n, iterable):
"Return first n items of the iterable as a list"
return list(islice(iterable, n))
def prepend(value, iterable):
"Prepend a single value in front of an iterable"
# prepend(1, [2, 3, 4]) --> 1 2 3 4
return chain([value], iterable)
def tabulate(function, start=0):
"Return function(0), function(1), ..."
return map(function, count(start))
def tail(n, iterable):
"Return an iterator over the last n items"
# tail(3, 'ABCDEFG') --> E F G
return iter(collections.deque(iterable, maxlen=n))
def consume(iterator, n=None):
"Advance the iterator n-steps ahead. If n is None, consume entirely."
# Use functions that consume iterators at C speed.
if n is None:
# feed the entire iterator into a zero-length deque
collections.deque(iterator, maxlen=0)
else:
# advance to the empty slice starting at position n
next(islice(iterator, n, n), None)
def nth(iterable, n, default=None):
"Returns the nth item or a default value"
return next(islice(iterable, n, None), default)
def all_equal(iterable):
"Returns True if all the elements are equal to each other"
g = groupby(iterable)
return next(g, True) and not next(g, False)
def quantify(iterable, pred=bool):
"Count how many times the predicate is True"
return sum(map(pred, iterable))
def ncycles(iterable, n):
"Returns the sequence elements n times"
return chain.from_iterable(repeat(tuple(iterable), n))
def batched(iterable, n):
"Batch data into tuples of length n. The last batch may be shorter."
# batched('ABCDEFG', 3) --> ABC DEF G
if n < 1:
raise ValueError('n must be at least one')
it = iter(iterable)
while batch := tuple(islice(it, n)):
yield batch
def grouper(iterable, n, *, incomplete='fill', fillvalue=None):
"Collect data into non-overlapping fixed-length chunks or blocks"
# grouper('ABCDEFG', 3, fillvalue='x') --> ABC DEF Gxx
# grouper('ABCDEFG', 3, incomplete='strict') --> ABC DEF ValueError
# grouper('ABCDEFG', 3, incomplete='ignore') --> ABC DEF
args = [iter(iterable)] * n
if incomplete == 'fill':
return zip_longest(*args, fillvalue=fillvalue)
if incomplete == 'strict':
return zip(*args, strict=True)
if incomplete == 'ignore':
return zip(*args)
else:
raise ValueError('Expected fill, strict, or ignore')
def sumprod(vec1, vec2):
"Compute a sum of products."
return sum(starmap(operator.mul, zip(vec1, vec2, strict=True)))
def sum_of_squares(it):
"Add up the squares of the input values."
# sum_of_squares([10, 20, 30]) -> 1400
return sumprod(*tee(it))
def transpose(it):
"Swap the rows and columns of the input."
# transpose([(1, 2, 3), (11, 22, 33)]) --> (1, 11) (2, 22) (3, 33)
return zip(*it, strict=True)
def matmul(m1, m2):
"Multiply two matrices."
# matmul([(7, 5), (3, 5)], [[2, 5], [7, 9]]) --> (49, 80), (41, 60)
n = len(m2[0])
return batched(starmap(sumprod, product(m1, transpose(m2))), n)
def convolve(signal, kernel):
# See: https://betterexplained.com/articles/intuitive-convolution/
# convolve(data, [0.25, 0.25, 0.25, 0.25]) --> Moving average (blur)
# convolve(data, [1, -1]) --> 1st finite difference (1st derivative)
# convolve(data, [1, -2, 1]) --> 2nd finite difference (2nd derivative)
kernel = tuple(kernel)[::-1]
n = len(kernel)
window = collections.deque([0], maxlen=n) * n
for x in chain(signal, repeat(0, n-1)):
window.append(x)
yield sumprod(kernel, window)
def polynomial_from_roots(roots):
"""Compute a polynomial's coefficients from its roots.
(x - 5) (x + 4) (x - 3) expands to: x³ -4x² -17x + 60
"""
# polynomial_from_roots([5, -4, 3]) --> [1, -4, -17, 60]
expansion = [1]
for r in roots:
expansion = convolve(expansion, (1, -r))
return list(expansion)
def polynomial_eval(coefficients, x):
"""Evaluate a polynomial at a specific value.
Computes with better numeric stability than Horner's method.
"""
# Evaluate x³ -4x² -17x + 60 at x = 2.5
# polynomial_eval([1, -4, -17, 60], x=2.5) --> 8.125
n = len(coefficients)
if n == 0:
return x * 0 # coerce zero to the type of x
powers = map(pow, repeat(x), reversed(range(n)))
return sumprod(coefficients, powers)
def iter_index(iterable, value, start=0):
"Return indices where a value occurs in a sequence or iterable."
# iter_index('AABCADEAF', 'A') --> 0 1 4 7
try:
seq_index = iterable.index
except AttributeError:
# Slow path for general iterables
it = islice(iterable, start, None)
i = start - 1
try:
while True:
yield (i := i + operator.indexOf(it, value) + 1)
except ValueError:
pass
else:
# Fast path for sequences
i = start - 1
try:
while True:
yield (i := seq_index(value, i+1))
except ValueError:
pass
def sieve(n):
"Primes less than n"
# sieve(30) --> 2 3 5 7 11 13 17 19 23 29
data = bytearray((0, 1)) * (n // 2)
data[:3] = 0, 0, 0
limit = math.isqrt(n) + 1
for p in compress(range(limit), data):
data[p*p : n : p+p] = bytes(len(range(p*p, n, p+p)))
data[2] = 1
return iter_index(data, 1) if n > 2 else iter([])
def factor(n):
"Prime factors of n."
# factor(99) --> 3 3 11
for prime in sieve(math.isqrt(n) + 1):
while True:
quotient, remainder = divmod(n, prime)
if remainder:
break
yield prime
n = quotient
if n == 1:
return
if n > 1:
yield n
def flatten(list_of_lists):
"Flatten one level of nesting"
return chain.from_iterable(list_of_lists)
def repeatfunc(func, times=None, *args):
"""Repeat calls to func with specified arguments.
Example: repeatfunc(random.random)
"""
if times is None:
return starmap(func, repeat(args))
return starmap(func, repeat(args, times))
def triplewise(iterable):
"Return overlapping triplets from an iterable"
# triplewise('ABCDEFG') --> ABC BCD CDE DEF EFG
for (a, _), (b, c) in pairwise(pairwise(iterable)):
yield a, b, c
def sliding_window(iterable, n):
# sliding_window('ABCDEFG', 4) --> ABCD BCDE CDEF DEFG
it = iter(iterable)
window = collections.deque(islice(it, n), maxlen=n)
if len(window) == n:
yield tuple(window)
for x in it:
window.append(x)
yield tuple(window)
def roundrobin(*iterables):
"roundrobin('ABC', 'D', 'EF') --> A D E B F C"
# Recipe credited to George Sakkis
num_active = len(iterables)
nexts = 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 = cycle(islice(nexts, num_active))
def partition(pred, iterable):
"Use a predicate to partition entries into false entries and true entries"
# partition(is_odd, range(10)) --> 0 2 4 6 8 and 1 3 5 7 9
t1, t2 = tee(iterable)
return filterfalse(pred, t1), filter(pred, t2)
def before_and_after(predicate, it):
""" Variant of takewhile() that allows complete
access to the remainder of the iterator.
>>> it = iter('ABCdEfGhI')
>>> all_upper, remainder = before_and_after(str.isupper, it)
>>> ''.join(all_upper)
'ABC'
>>> ''.join(remainder) # takewhile() would lose the 'd'
'dEfGhI'
Note that the first iterator must be fully
consumed before the second iterator can
generate valid results.
"""
it = iter(it)
transition = []
def true_iterator():
for elem in it:
if predicate(elem):
yield elem
else:
transition.append(elem)
return
def remainder_iterator():
yield from transition
yield from it
return true_iterator(), remainder_iterator()
def subslices(seq):
"Return all contiguous non-empty subslices of a sequence"
# subslices('ABCD') --> A AB ABC ABCD B BC BCD C CD D
slices = starmap(slice, combinations(range(len(seq) + 1), 2))
return map(operator.getitem, repeat(seq), slices)
def powerset(iterable):
"powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
s = list(iterable)
return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
def unique_everseen(iterable, key=None):
"List unique elements, preserving order. Remember all elements ever seen."
# unique_everseen('AAAABBBCCDAABBB') --> A B C D
# unique_everseen('ABBcCAD', str.lower) --> A B c D
seen = set()
if key is None:
for element in filterfalse(seen.__contains__, iterable):
seen.add(element)
yield element
# For order preserving deduplication,
# a faster but non-lazy solution is:
# yield from dict.fromkeys(iterable)
else:
for element in iterable:
k = key(element)
if k not in seen:
seen.add(k)
yield element
# For use cases that allow the last matching element to be returned,
# a faster but non-lazy solution is:
# t1, t2 = tee(iterable)
# yield from dict(zip(map(key, t1), t2)).values()
def unique_justseen(iterable, key=None):
"List unique elements, preserving order. Remember only the element just seen."
# 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), groupby(iterable, key)))
def iter_except(func, exception, first=None):
""" Call a function repeatedly until an exception is raised.
Converts a call-until-exception interface to an iterator interface.
Like builtins.iter(func, sentinel) but uses an exception instead
of a sentinel to end the loop.
Examples:
iter_except(functools.partial(heappop, h), IndexError) # priority queue iterator
iter_except(d.popitem, KeyError) # non-blocking dict iterator
iter_except(d.popleft, IndexError) # non-blocking deque iterator
iter_except(q.get_nowait, Queue.Empty) # loop over a producer Queue
iter_except(s.pop, KeyError) # non-blocking set iterator
"""
try:
if first is not None:
yield first() # For database APIs needing an initial cast to db.first()
while True:
yield func()
except exception:
pass
def first_true(iterable, default=False, pred=None):
"""Returns the first true value in the iterable.
If no true value is found, returns *default*
If *pred* is not None, returns the first item
for which pred(item) is true.
"""
# 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)
def nth_combination(iterable, r, index):
"Equivalent to list(combinations(iterable, r))[index]"
pool = tuple(iterable)
n = len(pool)
c = math.comb(n, r)
if index < 0:
index += c
if index < 0 or index >= c:
raise IndexError
result = []
while r:
c, n, r = c*r//n, n-1, r-1
while index >= c:
index -= c
c, n = c*(n-r)//n, n-1
result.append(pool[-1-n])
return tuple(result)
© 2001–2023 Python Software Foundation
Licensed under the PSF License.
https://docs.python.org/3.11/library/itertools.html