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