Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""Imported from the recipes section of the itertools documentation.
3All functions taken from the recipes section of the itertools library docs
4[1]_.
5Some backward-compatible usability improvements have been made.
7.. [1] http://docs.python.org/library/itertools.html#recipes
9"""
11import random
13from bisect import bisect_left, insort
14from collections import deque
15from contextlib import suppress
16from dataclasses import dataclass
17from functools import lru_cache, reduce
18from heapq import heappush, heappushpop
19from itertools import (
20 accumulate,
21 chain,
22 combinations,
23 compress,
24 count,
25 cycle,
26 filterfalse,
27 groupby,
28 islice,
29 pairwise as itertools_pairwise,
30 product,
31 repeat,
32 starmap,
33 takewhile,
34 tee,
35 zip_longest,
36)
37from math import prod, comb, isqrt, gcd
38from operator import mul, getitem, index as _index, is_, itemgetter, truediv
39from random import randrange, sample, choice, shuffle
40from sys import hexversion
42__all__ = [
43 'Stats',
44 'all_equal',
45 'batched',
46 'before_and_after',
47 'consume',
48 'convolve',
49 'dotproduct',
50 'first_true',
51 'factor',
52 'flatten',
53 'grouper',
54 'is_prime',
55 'iter_except',
56 'iter_index',
57 'loops',
58 'matmul',
59 'multinomial',
60 'ncycles',
61 'nth',
62 'nth_combination',
63 'padnone',
64 'pad_none',
65 'pairwise',
66 'partition',
67 'polynomial_eval',
68 'polynomial_from_roots',
69 'polynomial_derivative',
70 'powerset',
71 'prepend',
72 'quantify',
73 'reshape',
74 'random_combination_with_replacement',
75 'random_combination',
76 'random_derangement',
77 'random_permutation',
78 'random_product',
79 'repeatfunc',
80 'roundrobin',
81 'running_max',
82 'running_mean',
83 'running_median',
84 'running_min',
85 'running_statistics',
86 'sieve',
87 'sliding_window',
88 'subslices',
89 'sum_of_squares',
90 'tabulate',
91 'tail',
92 'take',
93 'totient',
94 'transpose',
95 'triplewise',
96 'unique',
97 'unique_everseen',
98 'unique_justseen',
99]
101_marker = object()
104# heapq max-heap functions are available for Python 3.14+
105try:
106 from heapq import heappush_max, heappushpop_max
107except ImportError: # pragma: no cover
108 _max_heap_available = False
109else: # pragma: no cover
110 _max_heap_available = True
113def take(n, iterable):
114 """Return first *n* items of the *iterable* as a list.
116 >>> take(3, range(10))
117 [0, 1, 2]
119 If there are fewer than *n* items in the iterable, all of them are
120 returned.
122 >>> take(10, range(3))
123 [0, 1, 2]
125 """
126 return list(islice(iterable, n))
129def tabulate(function, start=0):
130 """Return an iterator over the results of ``func(start)``,
131 ``func(start + 1)``, ``func(start + 2)``...
133 *func* should be a function that accepts one integer argument.
135 If *start* is not specified it defaults to 0. It will be incremented each
136 time the iterator is advanced.
138 >>> square = lambda x: x ** 2
139 >>> iterator = tabulate(square, -3)
140 >>> take(4, iterator)
141 [9, 4, 1, 0]
143 """
144 return map(function, count(start))
147def tail(n, iterable):
148 """Return an iterator over the last *n* items of *iterable*.
150 >>> t = tail(3, 'ABCDEFG')
151 >>> list(t)
152 ['E', 'F', 'G']
154 """
155 try:
156 size = len(iterable)
157 except TypeError:
158 return iter(deque(iterable, maxlen=n))
159 else:
160 return islice(iterable, max(0, size - n), None)
163def consume(iterator, n=None):
164 """Advance *iterable* by *n* steps. If *n* is ``None``, consume it
165 entirely.
167 Efficiently exhausts an iterator without returning values. Defaults to
168 consuming the whole iterator, but an optional second argument may be
169 provided to limit consumption.
171 >>> i = (x for x in range(10))
172 >>> next(i)
173 0
174 >>> consume(i, 3)
175 >>> next(i)
176 4
177 >>> consume(i)
178 >>> next(i)
179 Traceback (most recent call last):
180 File "<stdin>", line 1, in <module>
181 StopIteration
183 If the iterator has fewer items remaining than the provided limit, the
184 whole iterator will be consumed.
186 >>> i = (x for x in range(3))
187 >>> consume(i, 5)
188 >>> next(i)
189 Traceback (most recent call last):
190 File "<stdin>", line 1, in <module>
191 StopIteration
193 """
194 # Use functions that consume iterators at C speed.
195 if n is None:
196 # feed the entire iterator into a zero-length deque
197 deque(iterator, maxlen=0)
198 else:
199 # advance to the empty slice starting at position n
200 next(islice(iterator, n, n), None)
203def nth(iterable, n, default=None):
204 """Returns the nth item or a default value.
206 >>> l = range(10)
207 >>> nth(l, 3)
208 3
209 >>> nth(l, 20, "zebra")
210 'zebra'
212 """
213 return next(islice(iterable, n, None), default)
216def all_equal(iterable, key=None):
217 """
218 Returns ``True`` if all the elements are equal to each other.
220 >>> all_equal('aaaa')
221 True
222 >>> all_equal('aaab')
223 False
225 A function that accepts a single argument and returns a transformed version
226 of each input item can be specified with *key*:
228 >>> all_equal('AaaA', key=str.casefold)
229 True
230 >>> all_equal([1, 2, 3], key=lambda x: x < 10)
231 True
233 """
234 iterator = groupby(iterable, key)
235 for first in iterator:
236 for second in iterator:
237 return False
238 return True
239 return True
242def quantify(iterable, pred=bool):
243 """Return the how many times the predicate is true.
245 >>> quantify([True, False, True])
246 2
248 """
249 return sum(map(pred, iterable))
252def pad_none(iterable):
253 """Returns the sequence of elements and then returns ``None`` indefinitely.
255 >>> take(5, pad_none(range(3)))
256 [0, 1, 2, None, None]
258 Useful for emulating the behavior of the built-in :func:`map` function.
260 See also :func:`padded`.
262 """
263 return chain(iterable, repeat(None))
266padnone = pad_none
269def ncycles(iterable, n):
270 """Returns the sequence elements *n* times
272 >>> list(ncycles(["a", "b"], 3))
273 ['a', 'b', 'a', 'b', 'a', 'b']
275 """
276 return chain.from_iterable(repeat(tuple(iterable), n))
279def dotproduct(vec1, vec2):
280 """Returns the dot product of the two iterables.
282 >>> dotproduct([10, 15, 12], [0.65, 0.80, 1.25])
283 33.5
284 >>> 10 * 0.65 + 15 * 0.80 + 12 * 1.25
285 33.5
287 In Python 3.12 and later, use ``math.sumprod()`` instead.
288 """
289 return sum(map(mul, vec1, vec2))
292# math.sumprod is available for Python 3.12+
293try:
294 from math import sumprod as _sumprod
295except ImportError: # pragma: no cover
296 _sumprod = dotproduct
299def flatten(list_of_lists):
300 """Return an iterator flattening one level of nesting in a list of lists.
302 >>> list(flatten([[0, 1], [2, 3]]))
303 [0, 1, 2, 3]
305 See also :func:`collapse`, which can flatten multiple levels of nesting.
307 """
308 return chain.from_iterable(list_of_lists)
311def repeatfunc(function, times=None, *args):
312 """Call *function* with *args* repeatedly, returning an iterable over the
313 results.
315 If *times* is specified, the iterable will terminate after that many
316 repetitions:
318 >>> from operator import add
319 >>> times = 4
320 >>> args = 3, 5
321 >>> list(repeatfunc(add, times, *args))
322 [8, 8, 8, 8]
324 If *times* is ``None`` the iterable will not terminate:
326 >>> from random import randrange
327 >>> times = None
328 >>> args = 1, 11
329 >>> take(6, repeatfunc(randrange, times, *args)) # doctest:+SKIP
330 [2, 4, 8, 1, 8, 4]
332 """
333 if times is None:
334 return starmap(function, repeat(args))
335 return starmap(function, repeat(args, times))
338def pairwise(iterable):
339 """
340 Wrapper for :func:`itertools.pairwise`.
342 .. warning::
344 This function is deprecated as of version 11.0.0. It will be removed in a future
345 major release.
346 """
347 return itertools_pairwise(iterable)
350def grouper(iterable, n, incomplete='fill', fillvalue=None):
351 """Group elements from *iterable* into fixed-length groups of length *n*.
353 >>> list(grouper('ABCDEF', 3))
354 [('A', 'B', 'C'), ('D', 'E', 'F')]
356 The keyword arguments *incomplete* and *fillvalue* control what happens for
357 iterables whose length is not a multiple of *n*.
359 When *incomplete* is `'fill'`, the last group will contain instances of
360 *fillvalue*.
362 >>> list(grouper('ABCDEFG', 3, incomplete='fill', fillvalue='x'))
363 [('A', 'B', 'C'), ('D', 'E', 'F'), ('G', 'x', 'x')]
365 When *incomplete* is `'ignore'`, the last group will not be emitted.
367 >>> list(grouper('ABCDEFG', 3, incomplete='ignore', fillvalue='x'))
368 [('A', 'B', 'C'), ('D', 'E', 'F')]
370 When *incomplete* is `'strict'`, a `ValueError` will be raised.
372 >>> iterator = grouper('ABCDEFG', 3, incomplete='strict')
373 >>> list(iterator) # doctest: +IGNORE_EXCEPTION_DETAIL
374 Traceback (most recent call last):
375 ...
376 ValueError
378 """
379 iterators = [iter(iterable)] * n
380 match incomplete:
381 case 'fill':
382 return zip_longest(*iterators, fillvalue=fillvalue)
383 case 'strict':
384 return zip(*iterators, strict=True)
385 case 'ignore':
386 return zip(*iterators)
387 case _:
388 raise ValueError('Expected fill, strict, or ignore')
391def roundrobin(*iterables):
392 """Visit input iterables in a cycle until each is exhausted.
394 >>> list(roundrobin('ABC', 'D', 'EF'))
395 ['A', 'D', 'E', 'B', 'F', 'C']
397 This function produces the same output as :func:`interleave_longest`, but
398 may perform better for some inputs (in particular when the number of
399 iterables is small).
401 """
402 # Algorithm credited to George Sakkis
403 iterators = map(iter, iterables)
404 for num_active in range(len(iterables), 0, -1):
405 iterators = cycle(islice(iterators, num_active))
406 yield from map(next, iterators)
409def partition(pred, iterable):
410 """
411 Returns a 2-tuple of iterables derived from the input iterable.
412 The first yields the items that have ``pred(item) == False``.
413 The second yields the items that have ``pred(item) == True``.
415 >>> is_odd = lambda x: x % 2 != 0
416 >>> iterable = range(10)
417 >>> even_items, odd_items = partition(is_odd, iterable)
418 >>> list(even_items), list(odd_items)
419 ([0, 2, 4, 6, 8], [1, 3, 5, 7, 9])
421 If *pred* is None, :func:`bool` is used.
423 >>> iterable = [0, 1, False, True, '', ' ']
424 >>> false_items, true_items = partition(None, iterable)
425 >>> list(false_items), list(true_items)
426 ([0, False, ''], [1, True, ' '])
428 """
429 if pred is None:
430 pred = bool
431 iterator = iter(iterable)
433 false_queue = deque()
434 true_queue = deque()
436 def gen(queue):
437 while True:
438 while queue:
439 yield queue.popleft()
440 for value in iterator:
441 (true_queue if pred(value) else false_queue).append(value)
442 break
443 else:
444 return
446 return gen(false_queue), gen(true_queue)
449def powerset(iterable):
450 """Yields all possible subsets of the iterable.
452 >>> list(powerset([1, 2, 3]))
453 [(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
455 :func:`powerset` will operate on iterables that aren't :class:`set`
456 instances, so repeated elements in the input will produce repeated elements
457 in the output.
459 >>> seq = [1, 1, 0]
460 >>> list(powerset(seq))
461 [(), (1,), (1,), (0,), (1, 1), (1, 0), (1, 0), (1, 1, 0)]
463 For a variant that efficiently yields actual :class:`set` instances, see
464 :func:`powerset_of_sets`.
465 """
466 s = list(iterable)
467 return chain.from_iterable(combinations(s, r) for r in range(len(s) + 1))
470def unique_everseen(iterable, key=None):
471 """Yield unique elements, preserving order. Remember all elements ever seen.
473 >>> list(unique_everseen('AAAABBBCCDAABBB'))
474 ['A', 'B', 'C', 'D']
475 >>> list(unique_everseen('ABBCcAD', str.casefold))
476 ['A', 'B', 'C', 'D']
478 Raises ``TypeError`` for unhashable items.
480 Some unhashable objects can be converted to hashable objects
481 using the *key* parameter:
483 * For ``list`` objects, try ``key=tuple``.
484 * For ``set`` objects, try ``key=frozenset``.
485 * For ``dict`` objects, try ``key=lambda x: frozenset(x.items())``
486 or in Python 3.15 and later, set ``key=frozendict``.
488 Alternatively, consider the ``unique()`` itertool recipe. It sorts
489 the data and then uses equality to eliminate duplicates. Hashability
490 is not required.
492 """
493 seen = set()
494 if key is None:
495 for element in filterfalse(seen.__contains__, iterable):
496 seen.add(element)
497 yield element
498 else:
499 for element in iterable:
500 k = key(element)
501 if k not in seen:
502 seen.add(k)
503 yield element
506def unique_justseen(iterable, key=None):
507 """Yields elements in order, ignoring serial duplicates
509 >>> list(unique_justseen('AAAABBBCCDAABBB'))
510 ['A', 'B', 'C', 'D', 'A', 'B']
511 >>> list(unique_justseen('ABBCcAD', str.lower))
512 ['A', 'B', 'C', 'A', 'D']
514 """
515 if key is None:
516 return map(itemgetter(0), groupby(iterable))
518 return map(next, map(itemgetter(1), groupby(iterable, key)))
521def unique(iterable, key=None, reverse=False):
522 """Yields unique elements in sorted order.
524 >>> list(unique([[1, 2], [3, 4], [1, 2]]))
525 [[1, 2], [3, 4]]
527 *key* and *reverse* are passed to :func:`sorted`.
529 >>> list(unique('ABBcCAD', str.casefold))
530 ['A', 'B', 'c', 'D']
531 >>> list(unique('ABBcCAD', str.casefold, reverse=True))
532 ['D', 'c', 'B', 'A']
534 The elements in *iterable* need not be hashable, but they must be
535 comparable for sorting to work.
536 """
537 sequenced = sorted(iterable, key=key, reverse=reverse)
538 return unique_justseen(sequenced, key=key)
541def iter_except(function, exception, first=None):
542 """Yields results from a function repeatedly until an exception is raised.
544 Converts a call-until-exception interface to an iterator interface.
545 Like ``iter(function, sentinel)``, but uses an exception instead of a sentinel
546 to end the loop.
548 >>> l = [0, 1, 2]
549 >>> list(iter_except(l.pop, IndexError))
550 [2, 1, 0]
552 Multiple exceptions can be specified as a stopping condition:
554 >>> l = [1, 2, 3, '...', 4, 5, 6]
555 >>> list(iter_except(lambda: 1 + l.pop(), (IndexError, TypeError)))
556 [7, 6, 5]
557 >>> list(iter_except(lambda: 1 + l.pop(), (IndexError, TypeError)))
558 [4, 3, 2]
559 >>> list(iter_except(lambda: 1 + l.pop(), (IndexError, TypeError)))
560 []
562 """
563 with suppress(exception):
564 if first is not None:
565 yield first()
566 while True:
567 yield function()
570def first_true(iterable, default=None, pred=None):
571 """
572 Returns the first true value in the iterable.
574 If no true value is found, returns *default*
576 If *pred* is not None, returns the first item for which
577 ``pred(item) == True`` .
579 >>> first_true(range(10))
580 1
581 >>> first_true(range(10), pred=lambda x: x > 5)
582 6
583 >>> first_true(range(10), default='missing', pred=lambda x: x > 9)
584 'missing'
586 """
587 return next(filter(pred, iterable), default)
590def random_product(*iterables, repeat=1):
591 """Draw an item at random from each of the input iterables.
593 >>> random_product('abc', range(4), 'XYZ') # doctest:+SKIP
594 ('c', 3, 'Z')
596 If *repeat* is provided as a keyword argument, that many items will be
597 drawn from each iterable.
599 >>> random_product('abcd', range(4), repeat=2) # doctest:+SKIP
600 ('a', 2, 'd', 3)
602 This equivalent to taking a random selection from
603 ``itertools.product(*args, repeat=repeat)``.
605 """
606 pools = tuple(map(tuple, iterables)) * repeat
607 return tuple(map(choice, pools))
610def random_permutation(iterable, r=None):
611 """Return a random *r* length permutation of the elements in *iterable*.
613 If *r* is not specified or is ``None``, then *r* defaults to the length of
614 *iterable*.
616 >>> random_permutation(range(5)) # doctest:+SKIP
617 (3, 4, 0, 1, 2)
619 This equivalent to taking a random selection from
620 ``itertools.permutations(iterable, r)``.
622 """
623 pool = tuple(iterable)
624 r = len(pool) if r is None else r
625 return tuple(sample(pool, r))
628def random_combination(iterable, r):
629 """Return a random *r* length subsequence of the elements in *iterable*.
631 >>> random_combination(range(5), 3) # doctest:+SKIP
632 (2, 3, 4)
634 This equivalent to taking a random selection from
635 ``itertools.combinations(iterable, r)``.
637 """
638 pool = tuple(iterable)
639 n = len(pool)
640 indices = sorted(sample(range(n), r))
641 return tuple([pool[i] for i in indices])
644def random_combination_with_replacement(iterable, r):
645 """Return a random *r* length subsequence of elements in *iterable*,
646 allowing individual elements to be repeated.
648 >>> random_combination_with_replacement(range(3), 5) # doctest:+SKIP
649 (0, 0, 1, 2, 2)
651 This equivalent to taking a random selection from
652 ``itertools.combinations_with_replacement(iterable, r)``.
654 """
655 pool = tuple(iterable)
656 n = len(pool)
657 indices = sorted(randrange(n) for i in range(r))
658 return tuple([pool[i] for i in indices])
661def nth_combination(iterable, r, index):
662 """Equivalent to ``list(combinations(iterable, r))[index]``.
664 The subsequences of *iterable* that are of length *r* can be ordered
665 lexicographically. :func:`nth_combination` computes the subsequence at
666 sort position *index* directly, without computing the previous
667 subsequences.
669 >>> nth_combination(range(5), 3, 5)
670 (0, 3, 4)
672 ``ValueError`` will be raised If *r* is negative.
673 ``IndexError`` will be raised if the given *index* is invalid.
674 """
675 pool = tuple(iterable)
676 n = len(pool)
677 c = comb(n, r)
679 if index < 0:
680 index += c
681 if not 0 <= index < c:
682 raise IndexError
684 result = []
685 while r:
686 c, n, r = c * r // n, n - 1, r - 1
687 while index >= c:
688 index -= c
689 c, n = c * (n - r) // n, n - 1
690 result.append(pool[-1 - n])
692 return tuple(result)
695def prepend(value, iterable):
696 """Yield *value*, followed by the elements in *iterable*.
698 >>> value = '0'
699 >>> iterable = ['1', '2', '3']
700 >>> list(prepend(value, iterable))
701 ['0', '1', '2', '3']
703 To prepend multiple values, see :func:`itertools.chain`
704 or :func:`value_chain`.
706 """
707 return chain([value], iterable)
710def convolve(signal, kernel):
711 """Discrete linear convolution of two iterables.
712 Equivalent to polynomial multiplication.
714 For example, multiplying ``(x² -x - 20)`` by ``(x - 3)``
715 gives ``(x³ -4x² -17x + 60)``.
717 >>> list(convolve([1, -1, -20], [1, -3]))
718 [1, -4, -17, 60]
720 Examples of popular kinds of kernels:
722 * The kernel ``[0.25, 0.25, 0.25, 0.25]`` computes a moving average.
723 For image data, this blurs the image and reduces noise.
724 * The kernel ``[1/2, 0, -1/2]`` estimates the first derivative of
725 a function evaluated at evenly spaced inputs.
726 * The kernel ``[1, -2, 1]`` estimates the second derivative of a
727 function evaluated at evenly spaced inputs.
729 Convolutions are mathematically commutative; however, the inputs are
730 evaluated differently. The signal is consumed lazily and can be
731 infinite. The kernel is fully consumed before the calculations begin.
733 Supports all numeric types: int, float, complex, Decimal, Fraction.
735 References:
737 * Article: https://betterexplained.com/articles/intuitive-convolution/
738 * Video by 3Blue1Brown: https://www.youtube.com/watch?v=KuXjwB4LzSA
740 """
741 # This implementation comes from an older version of the itertools
742 # documentation. While the newer implementation is a bit clearer,
743 # this one was kept because the inlined window logic is faster
744 # and it avoids an unnecessary deque-to-tuple conversion.
745 kernel = tuple(kernel)[::-1]
746 n = len(kernel)
747 window = deque([0], maxlen=n) * n
748 for x in chain(signal, repeat(0, n - 1)):
749 window.append(x)
750 yield _sumprod(kernel, window)
753def before_and_after(predicate, it):
754 """A variant of :func:`takewhile` that allows complete access to the
755 remainder of the iterator.
757 >>> it = iter('ABCdEfGhI')
758 >>> all_upper, remainder = before_and_after(str.isupper, it)
759 >>> ''.join(all_upper)
760 'ABC'
761 >>> ''.join(remainder) # takewhile() would lose the 'd'
762 'dEfGhI'
764 Note that the first iterator must be fully consumed before the second
765 iterator can generate valid results.
766 """
767 trues, after = tee(it)
768 trues = compress(takewhile(predicate, trues), zip(after))
769 return trues, after
772def triplewise(iterable):
773 """Return overlapping triplets from *iterable*.
775 >>> list(triplewise('ABCDE'))
776 [('A', 'B', 'C'), ('B', 'C', 'D'), ('C', 'D', 'E')]
778 """
779 # This deviates from the itertools documentation recipe - see
780 # https://github.com/more-itertools/more-itertools/issues/889
781 t1, t2, t3 = tee(iterable, 3)
782 next(t3, None)
783 next(t3, None)
784 next(t2, None)
785 return zip(t1, t2, t3)
788def _sliding_window_islice(iterable, n):
789 # Fast path for small, non-zero values of n.
790 iterators = tee(iterable, n)
791 for i, iterator in enumerate(iterators):
792 next(islice(iterator, i, i), None)
793 return zip(*iterators)
796def _sliding_window_deque(iterable, n):
797 # Normal path for other values of n.
798 iterator = iter(iterable)
799 window = deque(islice(iterator, n - 1), maxlen=n)
800 for x in iterator:
801 window.append(x)
802 yield tuple(window)
805def sliding_window(iterable, n):
806 """Return a sliding window of width *n* over *iterable*.
808 >>> list(sliding_window(range(6), 4))
809 [(0, 1, 2, 3), (1, 2, 3, 4), (2, 3, 4, 5)]
811 If *iterable* has fewer than *n* items, then nothing is yielded:
813 >>> list(sliding_window(range(3), 4))
814 []
816 For a variant with more features, see :func:`windowed`.
817 """
818 if n > 20:
819 return _sliding_window_deque(iterable, n)
820 elif n > 2:
821 return _sliding_window_islice(iterable, n)
822 elif n == 2:
823 return pairwise(iterable)
824 elif n == 1:
825 return zip(iterable)
826 else:
827 raise ValueError(f'n should be at least one, not {n}')
830def subslices(iterable):
831 """Return all contiguous non-empty subslices of *iterable*.
833 >>> list(subslices('ABC'))
834 [['A'], ['A', 'B'], ['A', 'B', 'C'], ['B'], ['B', 'C'], ['C']]
836 This is similar to :func:`substrings`, but emits items in a different
837 order.
838 """
839 seq = list(iterable)
840 slices = starmap(slice, combinations(range(len(seq) + 1), 2))
841 return map(getitem, repeat(seq), slices)
844def polynomial_from_roots(roots):
845 """Compute a polynomial's coefficients from its roots.
847 >>> roots = [5, -4, 3] # (x - 5) * (x + 4) * (x - 3)
848 >>> polynomial_from_roots(roots) # x³ - 4 x² - 17 x + 60
849 [1, -4, -17, 60]
851 Note that polynomial coefficients are specified in descending power order.
853 Supports all numeric types: int, float, complex, Decimal, Fraction.
854 """
856 # This recipe differs from the one in itertools docs in that it
857 # applies list() after each call to convolve(). This avoids
858 # hitting stack limits with nested generators.
860 poly = [1]
861 for root in roots:
862 poly = list(convolve(poly, (1, -root)))
863 return poly
866def iter_index(iterable, value, start=0, stop=None):
867 """Yield the index of each place in *iterable* that *value* occurs,
868 beginning with index *start* and ending before index *stop*.
871 >>> list(iter_index('AABCADEAF', 'A'))
872 [0, 1, 4, 7]
873 >>> list(iter_index('AABCADEAF', 'A', 1)) # start index is inclusive
874 [1, 4, 7]
875 >>> list(iter_index('AABCADEAF', 'A', 1, 7)) # stop index is not inclusive
876 [1, 4]
878 The behavior for non-scalar *values* matches the built-in Python types.
880 >>> list(iter_index('ABCDABCD', 'AB'))
881 [0, 4]
882 >>> list(iter_index([0, 1, 2, 3, 0, 1, 2, 3], [0, 1]))
883 []
884 >>> list(iter_index([[0, 1], [2, 3], [0, 1], [2, 3]], [0, 1]))
885 [0, 2]
887 See :func:`locate` for a more general means of finding the indexes
888 associated with particular values.
890 """
891 seq_index = getattr(iterable, 'index', None)
892 if seq_index is None:
893 # Slow path for general iterables
894 iterator = islice(iterable, start, stop)
895 for i, element in enumerate(iterator, start):
896 if element is value or element == value:
897 yield i
898 else:
899 # Fast path for sequences
900 stop = len(iterable) if stop is None else stop
901 i = start - 1
902 with suppress(ValueError):
903 while True:
904 yield (i := seq_index(value, i + 1, stop))
907def sieve(n):
908 """Yield the primes less than n.
910 >>> list(sieve(30))
911 [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]
913 """
914 # This implementation comes from an older version of the itertools
915 # documentation. The newer implementation is easier to read but is
916 # less lazy.
917 if n > 2:
918 yield 2
919 start = 3
920 data = bytearray((0, 1)) * (n // 2)
921 for p in iter_index(data, 1, start, stop=isqrt(n) + 1):
922 yield from iter_index(data, 1, start, p * p)
923 data[p * p : n : p + p] = bytes(len(range(p * p, n, p + p)))
924 start = p * p
925 yield from iter_index(data, 1, start)
928def _batched(iterable, n, *, strict=False): # pragma: no cover
929 """Batch data into tuples of length *n*. If the number of items in
930 *iterable* is not divisible by *n*:
931 * The last batch will be shorter if *strict* is ``False``.
932 * :exc:`ValueError` will be raised if *strict* is ``True``.
934 >>> list(batched('ABCDEFG', 3))
935 [('A', 'B', 'C'), ('D', 'E', 'F'), ('G',)]
937 On Python 3.13 and above, this is an alias for :func:`itertools.batched`.
938 """
939 if n < 1:
940 raise ValueError('n must be at least one')
941 iterator = iter(iterable)
942 while batch := tuple(islice(iterator, n)):
943 if strict and len(batch) != n:
944 raise ValueError('batched(): incomplete batch')
945 yield batch
948if hexversion >= 0x30D00A2: # pragma: no cover
949 from itertools import batched as itertools_batched
951 def batched(iterable, n, *, strict=False):
952 return itertools_batched(iterable, n, strict=strict)
954 batched.__doc__ = _batched.__doc__
955else: # pragma: no cover
956 batched = _batched
959def transpose(matrix):
960 """Swap the rows and columns of the input matrix.
962 >>> list(transpose([(1, 2, 3), (11, 22, 33)]))
963 [(1, 11), (2, 22), (3, 33)]
965 The caller should ensure that the dimensions of the input are compatible.
966 If the input is empty, no output will be produced.
967 """
968 return zip(*matrix, strict=True)
971def _is_scalar(value, stringlike=(str, bytes)):
972 "Scalars are bytes, strings, and non-iterables."
973 try:
974 iter(value)
975 except TypeError:
976 return True
977 return isinstance(value, stringlike)
980def _flatten_tensor(tensor):
981 "Depth-first iterator over scalars in a tensor."
982 iterator = iter(tensor)
983 while True:
984 try:
985 value = next(iterator)
986 except StopIteration:
987 return iterator
988 iterator = chain((value,), iterator)
989 if _is_scalar(value):
990 return iterator
991 iterator = chain.from_iterable(iterator)
994def reshape(matrix, shape):
995 """Change the shape of a *matrix*.
997 If *shape* is an integer, the matrix must be two dimensional
998 and the shape is interpreted as the desired number of columns:
1000 >>> matrix = [(0, 1), (2, 3), (4, 5)]
1001 >>> cols = 3
1002 >>> list(reshape(matrix, cols))
1003 [(0, 1, 2), (3, 4, 5)]
1005 If *shape* is a tuple (or other iterable), the input matrix can have
1006 any number of dimensions. It will first be flattened and then rebuilt
1007 to the desired shape which can also be multidimensional:
1009 >>> matrix = [(0, 1), (2, 3), (4, 5)] # Start with a 3 x 2 matrix
1011 >>> list(reshape(matrix, (2, 3))) # Make a 2 x 3 matrix
1012 [(0, 1, 2), (3, 4, 5)]
1014 >>> list(reshape(matrix, (6,))) # Make a vector of length six
1015 [0, 1, 2, 3, 4, 5]
1017 >>> list(reshape(matrix, (2, 1, 3, 1))) # Make 2 x 1 x 3 x 1 tensor
1018 [(((0,), (1,), (2,)),), (((3,), (4,), (5,)),)]
1020 Each dimension is assumed to be uniform, either all arrays or all scalars.
1021 Flattening stops when the first value in a dimension is a scalar.
1022 Scalars are bytes, strings, and non-iterables.
1023 The reshape iterator stops when the requested shape is complete
1024 or when the input is exhausted, whichever comes first.
1026 """
1027 if isinstance(shape, int):
1028 return batched(chain.from_iterable(matrix), shape)
1029 first_dim, *dims = shape
1030 scalar_stream = _flatten_tensor(matrix)
1031 reshaped = reduce(batched, reversed(dims), scalar_stream)
1032 return islice(reshaped, first_dim)
1035def matmul(m1, m2):
1036 """Multiply two matrices.
1038 >>> list(matmul([(7, 5), (3, 5)], [(2, 5), (7, 9)]))
1039 [(49, 80), (41, 60)]
1041 The caller should ensure that the dimensions of the input matrices are
1042 compatible with each other.
1044 Supports all numeric types: int, float, complex, Decimal, Fraction.
1045 """
1046 n = len(m2[0])
1047 return batched(starmap(_sumprod, product(m1, transpose(m2))), n)
1050def _factor_pollard(n):
1051 # Return a factor of n using Pollard's rho algorithm.
1052 # Efficient when n is odd and composite.
1053 for b in range(1, n):
1054 x = y = 2
1055 d = 1
1056 while d == 1:
1057 x = (x * x + b) % n
1058 y = (y * y + b) % n
1059 y = (y * y + b) % n
1060 d = gcd(x - y, n)
1061 if d != n:
1062 return d
1063 raise ValueError('prime or under 5') # pragma: no cover
1066_primes_below_211 = tuple(sieve(211))
1069def factor(n):
1070 """Yield the prime factors of n.
1072 >>> list(factor(360))
1073 [2, 2, 2, 3, 3, 5]
1075 Finds small factors with trial division. Larger factors are
1076 either verified as prime with ``is_prime`` or split into
1077 smaller factors with Pollard's rho algorithm.
1078 """
1080 # Corner case reduction
1081 if n < 2:
1082 return
1084 # Trial division reduction
1085 for prime in _primes_below_211:
1086 while not n % prime:
1087 yield prime
1088 n //= prime
1090 # Pollard's rho reduction
1091 primes = []
1092 todo = [n] if n > 1 else []
1093 for n in todo:
1094 if n < 211**2 or is_prime(n):
1095 primes.append(n)
1096 else:
1097 fact = _factor_pollard(n)
1098 todo += (fact, n // fact)
1099 yield from sorted(primes)
1102def polynomial_eval(coefficients, x):
1103 """Evaluate a polynomial at a specific value.
1105 Computes with better numeric stability than Horner's method.
1107 Evaluate ``x^3 - 4 * x^2 - 17 * x + 60`` at ``x = 2.5``:
1109 >>> coefficients = [1, -4, -17, 60]
1110 >>> x = 2.5
1111 >>> polynomial_eval(coefficients, x)
1112 8.125
1114 Note that polynomial coefficients are specified in descending power order.
1116 Supports all numeric types: int, float, complex, Decimal, Fraction.
1117 """
1118 n = len(coefficients)
1119 if n == 0:
1120 return type(x)(0)
1121 powers = map(pow, repeat(x), reversed(range(n)))
1122 return _sumprod(coefficients, powers)
1125def sum_of_squares(iterable):
1126 """Return the sum of the squares of the input values.
1128 >>> sum_of_squares([10, 20, 30])
1129 1400
1131 Supports all numeric types: int, float, complex, Decimal, Fraction.
1132 """
1133 return _sumprod(*tee(iterable))
1136def polynomial_derivative(coefficients):
1137 """Compute the first derivative of a polynomial.
1139 Evaluate the derivative of ``x³ - 4 x² - 17 x + 60``:
1141 >>> coefficients = [1, -4, -17, 60]
1142 >>> derivative_coefficients = polynomial_derivative(coefficients)
1143 >>> derivative_coefficients
1144 [3, -8, -17]
1146 Note that polynomial coefficients are specified in descending power order.
1148 Supports all numeric types: int, float, complex, Decimal, Fraction.
1149 """
1150 n = len(coefficients)
1151 powers = reversed(range(1, n))
1152 return list(map(mul, coefficients, powers))
1155def totient(n):
1156 """Return the count of natural numbers up to *n* that are coprime with *n*.
1158 Euler's totient function φ(n) gives the number of totatives.
1159 Totative are integers k in the range 1 ≤ k ≤ n such that gcd(n, k) = 1.
1161 >>> n = 9
1162 >>> totient(n)
1163 6
1165 >>> totatives = [x for x in range(1, n) if gcd(n, x) == 1]
1166 >>> totatives
1167 [1, 2, 4, 5, 7, 8]
1168 >>> len(totatives)
1169 6
1171 Reference: https://en.wikipedia.org/wiki/Euler%27s_totient_function
1173 """
1174 for prime in set(factor(n)):
1175 n -= n // prime
1176 return n
1179# Miller–Rabin primality test: https://oeis.org/A014233
1180_perfect_tests = [
1181 (2047, (2,)),
1182 (9080191, (31, 73)),
1183 (4759123141, (2, 7, 61)),
1184 (1122004669633, (2, 13, 23, 1662803)),
1185 (2152302898747, (2, 3, 5, 7, 11)),
1186 (3474749660383, (2, 3, 5, 7, 11, 13)),
1187 (18446744073709551616, (2, 325, 9375, 28178, 450775, 9780504, 1795265022)),
1188 (
1189 3317044064679887385961981,
1190 (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41),
1191 ),
1192]
1195@lru_cache
1196def _shift_to_odd(n):
1197 'Return s, d such that 2**s * d == n'
1198 s = ((n - 1) ^ n).bit_length() - 1
1199 d = n >> s
1200 assert (1 << s) * d == n and d & 1 and s >= 0
1201 return s, d
1204def _strong_probable_prime(n, base):
1205 assert (n > 2) and (n & 1) and (2 <= base < n)
1207 s, d = _shift_to_odd(n - 1)
1209 x = pow(base, d, n)
1210 if x == 1 or x == n - 1:
1211 return True
1213 for _ in range(s - 1):
1214 x = x * x % n
1215 if x == n - 1:
1216 return True
1218 return False
1221# Separate instance of Random() that doesn't share state
1222# with the default user instance of Random().
1223_private_randrange = random.Random().randrange
1226def is_prime(n):
1227 """Return ``True`` if *n* is prime and ``False`` otherwise.
1229 Basic examples:
1231 >>> is_prime(37)
1232 True
1233 >>> is_prime(3 * 13)
1234 False
1235 >>> is_prime(18_446_744_073_709_551_557)
1236 True
1238 Find the next prime over one billion:
1240 >>> next(filter(is_prime, count(10**9)))
1241 1000000007
1243 Generate random primes up to 200 bits and up to 60 decimal digits:
1245 >>> from random import seed, randrange, getrandbits
1246 >>> seed(18675309)
1248 >>> next(filter(is_prime, map(getrandbits, repeat(200))))
1249 893303929355758292373272075469392561129886005037663238028407
1251 >>> next(filter(is_prime, map(randrange, repeat(10**60))))
1252 269638077304026462407872868003560484232362454342414618963649
1254 This function is exact for values of *n* below 10**24. For larger inputs,
1255 the probabilistic Miller-Rabin primality test has a less than 1 in 2**128
1256 chance of a false positive.
1257 """
1259 if n < 17:
1260 return n in {2, 3, 5, 7, 11, 13}
1262 if not (n & 1 and n % 3 and n % 5 and n % 7 and n % 11 and n % 13):
1263 return False
1265 for limit, bases in _perfect_tests:
1266 if n < limit:
1267 break
1268 else:
1269 bases = (_private_randrange(2, n - 1) for i in range(64))
1271 return all(_strong_probable_prime(n, base) for base in bases)
1274def loops(n):
1275 """Returns an iterable with *n* elements for efficient looping.
1276 Like ``range(n)`` but doesn't create integers.
1278 >>> i = 0
1279 >>> for _ in loops(5):
1280 ... i += 1
1281 >>> i
1282 5
1284 """
1285 return repeat(None, n)
1288def multinomial(*counts):
1289 """Number of distinct arrangements of a multiset.
1291 The expression ``multinomial(3, 4, 2)`` has several equivalent
1292 interpretations:
1294 * In the expansion of ``(a + b + c)⁹``, the coefficient of the
1295 ``a³b⁴c²`` term is 1260.
1297 * There are 1260 distinct ways to arrange 9 balls consisting of 3 reds, 4
1298 greens, and 2 blues.
1300 * There are 1260 unique ways to place 9 distinct objects into three bins
1301 with sizes 3, 4, and 2.
1303 The :func:`multinomial` function computes the length of
1304 :func:`distinct_permutations`. For example, there are 83,160 distinct
1305 anagrams of the word "abracadabra":
1307 >>> from more_itertools import distinct_permutations, ilen
1308 >>> ilen(distinct_permutations('abracadabra'))
1309 83160
1311 This can be computed directly from the letter counts, 5a 2b 2r 1c 1d:
1313 >>> from collections import Counter
1314 >>> list(Counter('abracadabra').values())
1315 [5, 2, 2, 1, 1]
1316 >>> multinomial(5, 2, 2, 1, 1)
1317 83160
1319 A binomial coefficient is a special case of multinomial where there are
1320 only two categories. For example, the number of ways to arrange 12 balls
1321 with 5 reds and 7 blues is ``multinomial(5, 7)`` or ``math.comb(12, 5)``.
1323 Likewise, factorial is a special case of multinomial where
1324 the multiplicities are all just 1 so that
1325 ``multinomial(1, 1, 1, 1, 1, 1, 1) == math.factorial(7)``.
1327 Reference: https://en.wikipedia.org/wiki/Multinomial_theorem
1329 """
1330 return prod(map(comb, accumulate(counts), counts))
1333def _running_median_minheap_and_maxheap(iterator): # pragma: no cover
1334 "Non-windowed running_median() for Python 3.14+"
1336 read = iterator.__next__
1337 lo = [] # max-heap
1338 hi = [] # min-heap (same size as or one smaller than lo)
1340 with suppress(StopIteration):
1341 while True:
1342 heappush_max(lo, heappushpop(hi, read()))
1343 yield lo[0]
1345 heappush(hi, heappushpop_max(lo, read()))
1346 yield (lo[0] + hi[0]) / 2
1349def _running_median_minheap_only(iterator): # pragma: no cover
1350 "Backport of non-windowed running_median() for Python 3.13 and prior."
1352 read = iterator.__next__
1353 lo = [] # max-heap (actually a minheap with negated values)
1354 hi = [] # min-heap (same size as or one smaller than lo)
1356 with suppress(StopIteration):
1357 while True:
1358 heappush(lo, -heappushpop(hi, read()))
1359 yield -lo[0]
1361 heappush(hi, -heappushpop(lo, -read()))
1362 yield (hi[0] - lo[0]) / 2
1365def _running_median_windowed(iterator, maxlen):
1366 "Yield median of values in a sliding window."
1368 window = deque()
1369 ordered = []
1371 for x in iterator:
1372 window.append(x)
1373 insort(ordered, x)
1375 if len(ordered) > maxlen:
1376 i = bisect_left(ordered, window.popleft())
1377 del ordered[i]
1379 n = len(ordered)
1380 m = n // 2
1381 yield ordered[m] if n & 1 else (ordered[m - 1] + ordered[m]) / 2
1384def running_median(iterable, *, maxlen=None):
1385 """Cumulative median of values seen so far or values in a sliding window.
1387 Set *maxlen* to a positive integer to specify the maximum size
1388 of the sliding window. The default of *None* is equivalent to
1389 an unbounded window.
1391 For example:
1393 >>> list(running_median([5.0, 9.0, 4.0, 12.0, 8.0, 9.0]))
1394 [5.0, 7.0, 5.0, 7.0, 8.0, 8.5]
1395 >>> list(running_median([5.0, 9.0, 4.0, 12.0, 8.0, 9.0], maxlen=3))
1396 [5.0, 7.0, 5.0, 9.0, 8.0, 9.0]
1398 Supports numeric types such as int, float, Decimal, and Fraction,
1399 but not complex numbers which are unorderable.
1401 On version Python 3.13 and prior, max-heaps are simulated with
1402 negative values. The negation causes Decimal inputs to apply context
1403 rounding, making the results slightly different than that obtained
1404 by statistics.median().
1405 """
1407 iterator = iter(iterable)
1409 if maxlen is not None:
1410 maxlen = _index(maxlen)
1411 if maxlen <= 0:
1412 raise ValueError('Window size should be positive')
1413 return _running_median_windowed(iterator, maxlen)
1415 if not _max_heap_available:
1416 return _running_median_minheap_only(iterator) # pragma: no cover
1418 return _running_median_minheap_and_maxheap(iterator) # pragma: no cover
1421def _windowed_running_mean(iterator, n):
1422 window = deque()
1423 running_sum = 0
1424 for value in iterator:
1425 window.append(value)
1426 running_sum += value
1427 if len(window) > n:
1428 running_sum -= window.popleft()
1429 yield running_sum / len(window)
1432def running_mean(iterable, *, maxlen=None):
1433 """Cumulative mean of values seen so far or values in a sliding window.
1435 Set *maxlen* to a positive integer to specify the maximum size
1436 of the sliding window. The default of *None* is equivalent to
1437 an unbounded window.
1439 For example:
1441 >>> list(running_mean([40, 30, 50, 46, 39, 44]))
1442 [40.0, 35.0, 40.0, 41.5, 41.0, 41.5]
1444 >>> list(running_mean([40, 30, 50, 46, 39, 44], maxlen=3))
1445 [40.0, 35.0, 40.0, 42.0, 45.0, 43.0]
1447 Supports numeric types such as int, float, complex, Decimal, and Fraction.
1449 No extra effort is made to reduce round-off errors for float inputs.
1450 So the results may be slightly different from `statistics.mean`.
1452 """
1454 iterator = iter(iterable)
1456 if maxlen is None:
1457 return map(truediv, accumulate(iterator), count(1))
1459 if maxlen <= 0:
1460 raise ValueError('Window size should be positive')
1462 return _windowed_running_mean(iterator, maxlen)
1465def _windowed_running_min(iterator, maxlen):
1466 sis = deque() # Strictly increasing subsequence
1467 for index, value in enumerate(iterator):
1468 if sis and sis[0][0] == index - maxlen:
1469 sis.popleft()
1470 while sis and not sis[-1][1] < value: # Remove non-increasing values
1471 sis.pop()
1472 sis.append((index, value)) # Most recent value at position -1
1473 yield sis[0][1] # Window minimum at position 0
1476def running_min(iterable, *, maxlen=None):
1477 """Smallest of values seen so far or values in a sliding window.
1479 Set *maxlen* to a positive integer to specify the maximum size
1480 of the sliding window. The default of *None* is equivalent to
1481 an unbounded window.
1483 For example:
1485 >>> list(running_min([4, 3, 7, 0, 8, 1, 6, 2, 9, 5]))
1486 [4, 3, 3, 0, 0, 0, 0, 0, 0, 0]
1488 >>> list(running_min([4, 3, 7, 0, 8, 1, 6, 2, 9, 5], maxlen=3))
1489 [4, 3, 3, 0, 0, 0, 1, 1, 2, 2]
1491 Supports numeric types such as int, float, Decimal, and Fraction,
1492 but not complex numbers which are unorderable.
1493 """
1495 iterator = iter(iterable)
1497 if maxlen is None:
1498 return accumulate(iterator, func=min)
1500 if maxlen <= 0:
1501 raise ValueError('Window size should be positive')
1503 return _windowed_running_min(iterator, maxlen)
1506def _windowed_running_max(iterator, maxlen):
1507 sds = deque() # Strictly decreasing subsequence
1508 for index, value in enumerate(iterator):
1509 if sds and sds[0][0] == index - maxlen:
1510 sds.popleft()
1511 while sds and not sds[-1][1] > value: # Remove non-decreasing values
1512 sds.pop()
1513 sds.append((index, value)) # Most recent value at position -1
1514 yield sds[0][1] # Window maximum at position 0
1517def running_max(iterable, *, maxlen=None):
1518 """Largest of values seen so far or values in a sliding window.
1520 Set *maxlen* to a positive integer to specify the maximum size
1521 of the sliding window. The default of *None* is equivalent to
1522 an unbounded window.
1524 For example:
1526 >>> list(running_max([4, 3, 7, 0, 8, 1, 6, 2, 9, 5]))
1527 [4, 4, 7, 7, 8, 8, 8, 8, 9, 9]
1529 >>> list(running_max([4, 3, 7, 0, 8, 1, 6, 2, 9, 5], maxlen=3))
1530 [4, 4, 7, 7, 8, 8, 8, 6, 9, 9]
1532 Supports numeric types such as int, float, Decimal, and Fraction,
1533 but not complex numbers which are unorderable.
1534 """
1536 iterator = iter(iterable)
1538 if maxlen is None:
1539 return accumulate(iterator, func=max)
1541 if maxlen <= 0:
1542 raise ValueError('Window size should be positive')
1544 return _windowed_running_max(iterator, maxlen)
1547@dataclass(frozen=True, slots=True)
1548class Stats:
1549 size: int
1550 minimum: float
1551 median: float
1552 maximum: float
1553 mean: float
1556def running_statistics(iterable, *, maxlen=None):
1557 """Statistics for values seen so far or values in a sliding window.
1559 Set *maxlen* to a positive integer to specify the maximum size
1560 of the sliding window. The default of *None* is equivalent to
1561 an unbounded window.
1563 Yields instances of a ``Stats`` dataclass with fields for the dataset *size*,
1564 *minimum* value, *median* value, *maximum* value, and the arithmetic *mean*.
1566 Supports numeric types such as int, float, Decimal, and Fraction,
1567 but not complex numbers which are unorderable.
1568 """
1570 # fmt: off
1571 t0, t1, t2, t3 = tee(iterable, 4)
1572 return map(
1573 Stats,
1574 count(1) if maxlen is None else chain(range(1, maxlen), repeat(maxlen)),
1575 running_min(t0, maxlen=maxlen),
1576 running_median(t1, maxlen=maxlen),
1577 running_max(t2, maxlen=maxlen),
1578 running_mean(t3, maxlen=maxlen),
1579 )
1580 # fmt: on
1583def random_derangement(iterable):
1584 """Return a random derangement of elements in the iterable.
1586 Equivalent to but much faster than ``choice(list(derangements(iterable)))``.
1588 """
1589 seq = tuple(iterable)
1590 if len(seq) < 2:
1591 if len(seq) == 0:
1592 return ()
1593 raise IndexError('No derangments to choose from')
1594 perm = list(range(len(seq)))
1595 start = tuple(perm)
1596 while True:
1597 shuffle(perm)
1598 if not any(map(is_, start, perm)):
1599 return itemgetter(*perm)(seq)