codekingpro/portable-devtools
114k
1"""Imported from the recipes section of the itertools documentation.2 3All functions taken from the recipes section of the itertools library docs4[1]_.5Some backward-compatible usability improvements have been made.6 7.. [1] http://docs.python.org/library/itertools.html#recipes8 9"""10import math11import operator12 13from collections import deque14from collections.abc import Sized15from functools import partial, reduce16from itertools import (17 chain,18 combinations,19 compress,20 count,21 cycle,22 groupby,23 islice,24 product,25 repeat,26 starmap,27 tee,28 zip_longest,29)30from random import randrange, sample, choice31from sys import hexversion32 33__all__ = [34 'all_equal',35 'batched',36 'before_and_after',37 'consume',38 'convolve',39 'dotproduct',40 'first_true',41 'factor',42 'flatten',43 'grouper',44 'iter_except',45 'iter_index',46 'matmul',47 'ncycles',48 'nth',49 'nth_combination',50 'padnone',51 'pad_none',52 'pairwise',53 'partition',54 'polynomial_eval',55 'polynomial_from_roots',56 'polynomial_derivative',57 'powerset',58 'prepend',59 'quantify',60 'reshape',61 'random_combination_with_replacement',62 'random_combination',63 'random_permutation',64 'random_product',65 'repeatfunc',66 'roundrobin',67 'sieve',68 'sliding_window',69 'subslices',70 'sum_of_squares',71 'tabulate',72 'tail',73 'take',74 'totient',75 'transpose',76 'triplewise',77 'unique_everseen',78 'unique_justseen',79]80 81_marker = object()82 83 84# zip with strict is available for Python 3.10+85try:86 zip(strict=True)87except TypeError:88 _zip_strict = zip89else:90 _zip_strict = partial(zip, strict=True)91 92# math.sumprod is available for Python 3.12+93_sumprod = getattr(math, 'sumprod', lambda x, y: dotproduct(x, y))94 95 96def take(n, iterable):97 """Return first *n* items of the iterable as a list.98 99 >>> take(3, range(10))100 [0, 1, 2]101 102 If there are fewer than *n* items in the iterable, all of them are103 returned.104 105 >>> take(10, range(3))106 [0, 1, 2]107 108 """109 return list(islice(iterable, n))110 111 112def tabulate(function, start=0):113 """Return an iterator over the results of ``func(start)``,114 ``func(start + 1)``, ``func(start + 2)``...115 116 *func* should be a function that accepts one integer argument.117 118 If *start* is not specified it defaults to 0. It will be incremented each119 time the iterator is advanced.120 121 >>> square = lambda x: x ** 2122 >>> iterator = tabulate(square, -3)123 >>> take(4, iterator)124 [9, 4, 1, 0]125 126 """127 return map(function, count(start))128 129 130def tail(n, iterable):131 """Return an iterator over the last *n* items of *iterable*.132 133 >>> t = tail(3, 'ABCDEFG')134 >>> list(t)135 ['E', 'F', 'G']136 137 """138 # If the given iterable has a length, then we can use islice to get its139 # final elements. Note that if the iterable is not actually Iterable,140 # either islice or deque will throw a TypeError. This is why we don't141 # check if it is Iterable.142 if isinstance(iterable, Sized):143 yield from islice(iterable, max(0, len(iterable) - n), None)144 else:145 yield from iter(deque(iterable, maxlen=n))146 147 148def consume(iterator, n=None):149 """Advance *iterable* by *n* steps. If *n* is ``None``, consume it150 entirely.151 152 Efficiently exhausts an iterator without returning values. Defaults to153 consuming the whole iterator, but an optional second argument may be154 provided to limit consumption.155 156 >>> i = (x for x in range(10))157 >>> next(i)158 0159 >>> consume(i, 3)160 >>> next(i)161 4162 >>> consume(i)163 >>> next(i)164 Traceback (most recent call last):165 File "<stdin>", line 1, in <module>166 StopIteration167 168 If the iterator has fewer items remaining than the provided limit, the169 whole iterator will be consumed.170 171 >>> i = (x for x in range(3))172 >>> consume(i, 5)173 >>> next(i)174 Traceback (most recent call last):175 File "<stdin>", line 1, in <module>176 StopIteration177 178 """179 # Use functions that consume iterators at C speed.180 if n is None:181 # feed the entire iterator into a zero-length deque182 deque(iterator, maxlen=0)183 else:184 # advance to the empty slice starting at position n185 next(islice(iterator, n, n), None)186 187 188def nth(iterable, n, default=None):189 """Returns the nth item or a default value.190 191 >>> l = range(10)192 >>> nth(l, 3)193 3194 >>> nth(l, 20, "zebra")195 'zebra'196 197 """198 return next(islice(iterable, n, None), default)199 200 201def all_equal(iterable):202 """203 Returns ``True`` if all the elements are equal to each other.204 205 >>> all_equal('aaaa')206 True207 >>> all_equal('aaab')208 False209 210 """211 g = groupby(iterable)212 return next(g, True) and not next(g, False)213 214 215def quantify(iterable, pred=bool):216 """Return the how many times the predicate is true.217 218 >>> quantify([True, False, True])219 2220 221 """222 return sum(map(pred, iterable))223 224 225def pad_none(iterable):226 """Returns the sequence of elements and then returns ``None`` indefinitely.227 228 >>> take(5, pad_none(range(3)))229 [0, 1, 2, None, None]230 231 Useful for emulating the behavior of the built-in :func:`map` function.232 233 See also :func:`padded`.234 235 """236 return chain(iterable, repeat(None))237 238 239padnone = pad_none240 241 242def ncycles(iterable, n):243 """Returns the sequence elements *n* times244 245 >>> list(ncycles(["a", "b"], 3))246 ['a', 'b', 'a', 'b', 'a', 'b']247 248 """249 return chain.from_iterable(repeat(tuple(iterable), n))250 251 252def dotproduct(vec1, vec2):253 """Returns the dot product of the two iterables.254 255 >>> dotproduct([10, 10], [20, 20])256 400257 258 """259 return sum(map(operator.mul, vec1, vec2))260 261 262def flatten(listOfLists):263 """Return an iterator flattening one level of nesting in a list of lists.264 265 >>> list(flatten([[0, 1], [2, 3]]))266 [0, 1, 2, 3]267 268 See also :func:`collapse`, which can flatten multiple levels of nesting.269 270 """271 return chain.from_iterable(listOfLists)272 273 274def repeatfunc(func, times=None, *args):275 """Call *func* with *args* repeatedly, returning an iterable over the276 results.277 278 If *times* is specified, the iterable will terminate after that many279 repetitions:280 281 >>> from operator import add282 >>> times = 4283 >>> args = 3, 5284 >>> list(repeatfunc(add, times, *args))285 [8, 8, 8, 8]286 287 If *times* is ``None`` the iterable will not terminate:288 289 >>> from random import randrange290 >>> times = None291 >>> args = 1, 11292 >>> take(6, repeatfunc(randrange, times, *args)) # doctest:+SKIP293 [2, 4, 8, 1, 8, 4]294 295 """296 if times is None:297 return starmap(func, repeat(args))298 return starmap(func, repeat(args, times))299 300 301def _pairwise(iterable):302 """Returns an iterator of paired items, overlapping, from the original303 304 >>> take(4, pairwise(count()))305 [(0, 1), (1, 2), (2, 3), (3, 4)]306 307 On Python 3.10 and above, this is an alias for :func:`itertools.pairwise`.308 309 """310 a, b = tee(iterable)311 next(b, None)312 return zip(a, b)313 314 315try:316 from itertools import pairwise as itertools_pairwise317except ImportError:318 pairwise = _pairwise319else:320 321 def pairwise(iterable):322 return itertools_pairwise(iterable)323 324 pairwise.__doc__ = _pairwise.__doc__325 326 327class UnequalIterablesError(ValueError):328 def __init__(self, details=None):329 msg = 'Iterables have different lengths'330 if details is not None:331 msg += (': index 0 has length {}; index {} has length {}').format(332 *details333 )334 335 super().__init__(msg)336 337 338def _zip_equal_generator(iterables):339 for combo in zip_longest(*iterables, fillvalue=_marker):340 for val in combo:341 if val is _marker:342 raise UnequalIterablesError()343 yield combo344 345 346def _zip_equal(*iterables):347 # Check whether the iterables are all the same size.348 try:349 first_size = len(iterables[0])350 for i, it in enumerate(iterables[1:], 1):351 size = len(it)352 if size != first_size:353 raise UnequalIterablesError(details=(first_size, i, size))354 # All sizes are equal, we can use the built-in zip.355 return zip(*iterables)356 # If any one of the iterables didn't have a length, start reading357 # them until one runs out.358 except TypeError:359 return _zip_equal_generator(iterables)360 361 362def grouper(iterable, n, incomplete='fill', fillvalue=None):363 """Group elements from *iterable* into fixed-length groups of length *n*.364 365 >>> list(grouper('ABCDEF', 3))366 [('A', 'B', 'C'), ('D', 'E', 'F')]367 368 The keyword arguments *incomplete* and *fillvalue* control what happens for369 iterables whose length is not a multiple of *n*.370 371 When *incomplete* is `'fill'`, the last group will contain instances of372 *fillvalue*.373 374 >>> list(grouper('ABCDEFG', 3, incomplete='fill', fillvalue='x'))375 [('A', 'B', 'C'), ('D', 'E', 'F'), ('G', 'x', 'x')]376 377 When *incomplete* is `'ignore'`, the last group will not be emitted.378 379 >>> list(grouper('ABCDEFG', 3, incomplete='ignore', fillvalue='x'))380 [('A', 'B', 'C'), ('D', 'E', 'F')]381 382 When *incomplete* is `'strict'`, a subclass of `ValueError` will be raised.383 384 >>> it = grouper('ABCDEFG', 3, incomplete='strict')385 >>> list(it) # doctest: +IGNORE_EXCEPTION_DETAIL386 Traceback (most recent call last):387 ...388 UnequalIterablesError389 390 """391 args = [iter(iterable)] * n392 if incomplete == 'fill':393 return zip_longest(*args, fillvalue=fillvalue)394 if incomplete == 'strict':395 return _zip_equal(*args)396 if incomplete == 'ignore':397 return zip(*args)398 else:399 raise ValueError('Expected fill, strict, or ignore')400 401 402def roundrobin(*iterables):403 """Yields an item from each iterable, alternating between them.404 405 >>> list(roundrobin('ABC', 'D', 'EF'))406 ['A', 'D', 'E', 'B', 'F', 'C']407 408 This function produces the same output as :func:`interleave_longest`, but409 may perform better for some inputs (in particular when the number of410 iterables is small).411 412 """413 # Recipe credited to George Sakkis414 pending = len(iterables)415 nexts = cycle(iter(it).__next__ for it in iterables)416 while pending:417 try:418 for next in nexts:419 yield next()420 except StopIteration:421 pending -= 1422 nexts = cycle(islice(nexts, pending))423 424 425def partition(pred, iterable):426 """427 Returns a 2-tuple of iterables derived from the input iterable.428 The first yields the items that have ``pred(item) == False``.429 The second yields the items that have ``pred(item) == True``.430 431 >>> is_odd = lambda x: x % 2 != 0432 >>> iterable = range(10)433 >>> even_items, odd_items = partition(is_odd, iterable)434 >>> list(even_items), list(odd_items)435 ([0, 2, 4, 6, 8], [1, 3, 5, 7, 9])436 437 If *pred* is None, :func:`bool` is used.438 439 >>> iterable = [0, 1, False, True, '', ' ']440 >>> false_items, true_items = partition(None, iterable)441 >>> list(false_items), list(true_items)442 ([0, False, ''], [1, True, ' '])443 444 """445 if pred is None:446 pred = bool447 448 t1, t2, p = tee(iterable, 3)449 p1, p2 = tee(map(pred, p))450 return (compress(t1, map(operator.not_, p1)), compress(t2, p2))451 452 453def powerset(iterable):454 """Yields all possible subsets of the iterable.455 456 >>> list(powerset([1, 2, 3]))457 [(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]458 459 :func:`powerset` will operate on iterables that aren't :class:`set`460 instances, so repeated elements in the input will produce repeated elements461 in the output. Use :func:`unique_everseen` on the input to avoid generating462 duplicates:463 464 >>> seq = [1, 1, 0]465 >>> list(powerset(seq))466 [(), (1,), (1,), (0,), (1, 1), (1, 0), (1, 0), (1, 1, 0)]467 >>> from more_itertools import unique_everseen468 >>> list(powerset(unique_everseen(seq)))469 [(), (1,), (0,), (1, 0)]470 471 """472 s = list(iterable)473 return chain.from_iterable(combinations(s, r) for r in range(len(s) + 1))474 475 476def unique_everseen(iterable, key=None):477 """478 Yield unique elements, preserving order.479 480 >>> list(unique_everseen('AAAABBBCCDAABBB'))481 ['A', 'B', 'C', 'D']482 >>> list(unique_everseen('ABBCcAD', str.lower))483 ['A', 'B', 'C', 'D']484 485 Sequences with a mix of hashable and unhashable items can be used.486 The function will be slower (i.e., `O(n^2)`) for unhashable items.487 488 Remember that ``list`` objects are unhashable - you can use the *key*489 parameter to transform the list to a tuple (which is hashable) to490 avoid a slowdown.491 492 >>> iterable = ([1, 2], [2, 3], [1, 2])493 >>> list(unique_everseen(iterable)) # Slow494 [[1, 2], [2, 3]]495 >>> list(unique_everseen(iterable, key=tuple)) # Faster496 [[1, 2], [2, 3]]497 498 Similarly, you may want to convert unhashable ``set`` objects with499 ``key=frozenset``. For ``dict`` objects,500 ``key=lambda x: frozenset(x.items())`` can be used.501 502 """503 seenset = set()504 seenset_add = seenset.add505 seenlist = []506 seenlist_add = seenlist.append507 use_key = key is not None508 509 for element in iterable:510 k = key(element) if use_key else element511 try:512 if k not in seenset:513 seenset_add(k)514 yield element515 except TypeError:516 if k not in seenlist:517 seenlist_add(k)518 yield element519 520 521def unique_justseen(iterable, key=None):522 """Yields elements in order, ignoring serial duplicates523 524 >>> list(unique_justseen('AAAABBBCCDAABBB'))525 ['A', 'B', 'C', 'D', 'A', 'B']526 >>> list(unique_justseen('ABBCcAD', str.lower))527 ['A', 'B', 'C', 'A', 'D']528 529 """530 if key is None:531 return map(operator.itemgetter(0), groupby(iterable))532 533 return map(next, map(operator.itemgetter(1), groupby(iterable, key)))534 535 536def iter_except(func, exception, first=None):537 """Yields results from a function repeatedly until an exception is raised.538 539 Converts a call-until-exception interface to an iterator interface.540 Like ``iter(func, sentinel)``, but uses an exception instead of a sentinel541 to end the loop.542 543 >>> l = [0, 1, 2]544 >>> list(iter_except(l.pop, IndexError))545 [2, 1, 0]546 547 Multiple exceptions can be specified as a stopping condition:548 549 >>> l = [1, 2, 3, '...', 4, 5, 6]550 >>> list(iter_except(lambda: 1 + l.pop(), (IndexError, TypeError)))551 [7, 6, 5]552 >>> list(iter_except(lambda: 1 + l.pop(), (IndexError, TypeError)))553 [4, 3, 2]554 >>> list(iter_except(lambda: 1 + l.pop(), (IndexError, TypeError)))555 []556 557 """558 try:559 if first is not None:560 yield first()561 while 1:562 yield func()563 except exception:564 pass565 566 567def first_true(iterable, default=None, pred=None):568 """569 Returns the first true value in the iterable.570 571 If no true value is found, returns *default*572 573 If *pred* is not None, returns the first item for which574 ``pred(item) == True`` .575 576 >>> first_true(range(10))577 1578 >>> first_true(range(10), pred=lambda x: x > 5)579 6580 >>> first_true(range(10), default='missing', pred=lambda x: x > 9)581 'missing'582 583 """584 return next(filter(pred, iterable), default)585 586 587def random_product(*args, repeat=1):588 """Draw an item at random from each of the input iterables.589 590 >>> random_product('abc', range(4), 'XYZ') # doctest:+SKIP591 ('c', 3, 'Z')592 593 If *repeat* is provided as a keyword argument, that many items will be594 drawn from each iterable.595 596 >>> random_product('abcd', range(4), repeat=2) # doctest:+SKIP597 ('a', 2, 'd', 3)598 599 This equivalent to taking a random selection from600 ``itertools.product(*args, **kwarg)``.601 602 """603 pools = [tuple(pool) for pool in args] * repeat604 return tuple(choice(pool) for pool in pools)605 606 607def random_permutation(iterable, r=None):608 """Return a random *r* length permutation of the elements in *iterable*.609 610 If *r* is not specified or is ``None``, then *r* defaults to the length of611 *iterable*.612 613 >>> random_permutation(range(5)) # doctest:+SKIP614 (3, 4, 0, 1, 2)615 616 This equivalent to taking a random selection from617 ``itertools.permutations(iterable, r)``.618 619 """620 pool = tuple(iterable)621 r = len(pool) if r is None else r622 return tuple(sample(pool, r))623 624 625def random_combination(iterable, r):626 """Return a random *r* length subsequence of the elements in *iterable*.627 628 >>> random_combination(range(5), 3) # doctest:+SKIP629 (2, 3, 4)630 631 This equivalent to taking a random selection from632 ``itertools.combinations(iterable, r)``.633 634 """635 pool = tuple(iterable)636 n = len(pool)637 indices = sorted(sample(range(n), r))638 return tuple(pool[i] for i in indices)639 640 641def random_combination_with_replacement(iterable, r):642 """Return a random *r* length subsequence of elements in *iterable*,643 allowing individual elements to be repeated.644 645 >>> random_combination_with_replacement(range(3), 5) # doctest:+SKIP646 (0, 0, 1, 2, 2)647 648 This equivalent to taking a random selection from649 ``itertools.combinations_with_replacement(iterable, r)``.650 651 """652 pool = tuple(iterable)653 n = len(pool)654 indices = sorted(randrange(n) for i in range(r))655 return tuple(pool[i] for i in indices)656 657 658def nth_combination(iterable, r, index):659 """Equivalent to ``list(combinations(iterable, r))[index]``.660 661 The subsequences of *iterable* that are of length *r* can be ordered662 lexicographically. :func:`nth_combination` computes the subsequence at663 sort position *index* directly, without computing the previous664 subsequences.665 666 >>> nth_combination(range(5), 3, 5)667 (0, 3, 4)668 669 ``ValueError`` will be raised If *r* is negative or greater than the length670 of *iterable*.671 ``IndexError`` will be raised if the given *index* is invalid.672 """673 pool = tuple(iterable)674 n = len(pool)675 if (r < 0) or (r > n):676 raise ValueError677 678 c = 1679 k = min(r, n - r)680 for i in range(1, k + 1):681 c = c * (n - k + i) // i682 683 if index < 0:684 index += c685 686 if (index < 0) or (index >= c):687 raise IndexError688 689 result = []690 while r:691 c, n, r = c * r // n, n - 1, r - 1692 while index >= c:693 index -= c694 c, n = c * (n - r) // n, n - 1695 result.append(pool[-1 - n])696 697 return tuple(result)698 699 700def prepend(value, iterator):701 """Yield *value*, followed by the elements in *iterator*.702 703 >>> value = '0'704 >>> iterator = ['1', '2', '3']705 >>> list(prepend(value, iterator))706 ['0', '1', '2', '3']707 708 To prepend multiple values, see :func:`itertools.chain`709 or :func:`value_chain`.710 711 """712 return chain([value], iterator)713 714 715def convolve(signal, kernel):716 """Convolve the iterable *signal* with the iterable *kernel*.717 718 >>> signal = (1, 2, 3, 4, 5)719 >>> kernel = [3, 2, 1]720 >>> list(convolve(signal, kernel))721 [3, 8, 14, 20, 26, 14, 5]722 723 Note: the input arguments are not interchangeable, as the *kernel*724 is immediately consumed and stored.725 726 """727 # This implementation intentionally doesn't match the one in the itertools728 # documentation.729 kernel = tuple(kernel)[::-1]730 n = len(kernel)731 window = deque([0], maxlen=n) * n732 for x in chain(signal, repeat(0, n - 1)):733 window.append(x)734 yield _sumprod(kernel, window)735 736 737def before_and_after(predicate, it):738 """A variant of :func:`takewhile` that allows complete access to the739 remainder of the iterator.740 741 >>> it = iter('ABCdEfGhI')742 >>> all_upper, remainder = before_and_after(str.isupper, it)743 >>> ''.join(all_upper)744 'ABC'745 >>> ''.join(remainder) # takewhile() would lose the 'd'746 'dEfGhI'747 748 Note that the first iterator must be fully consumed before the second749 iterator can generate valid results.750 """751 it = iter(it)752 transition = []753 754 def true_iterator():755 for elem in it:756 if predicate(elem):757 yield elem758 else:759 transition.append(elem)760 return761 762 # Note: this is different from itertools recipes to allow nesting763 # before_and_after remainders into before_and_after again. See tests764 # for an example.765 remainder_iterator = chain(transition, it)766 767 return true_iterator(), remainder_iterator768 769 770def triplewise(iterable):771 """Return overlapping triplets from *iterable*.772 773 >>> list(triplewise('ABCDE'))774 [('A', 'B', 'C'), ('B', 'C', 'D'), ('C', 'D', 'E')]775 776 """777 for (a, _), (b, c) in pairwise(pairwise(iterable)):778 yield a, b, c779 780 781def sliding_window(iterable, n):782 """Return a sliding window of width *n* over *iterable*.783 784 >>> list(sliding_window(range(6), 4))785 [(0, 1, 2, 3), (1, 2, 3, 4), (2, 3, 4, 5)]786 787 If *iterable* has fewer than *n* items, then nothing is yielded:788 789 >>> list(sliding_window(range(3), 4))790 []791 792 For a variant with more features, see :func:`windowed`.793 """794 it = iter(iterable)795 window = deque(islice(it, n - 1), maxlen=n)796 for x in it:797 window.append(x)798 yield tuple(window)799 800 801def subslices(iterable):802 """Return all contiguous non-empty subslices of *iterable*.803 804 >>> list(subslices('ABC'))805 [['A'], ['A', 'B'], ['A', 'B', 'C'], ['B'], ['B', 'C'], ['C']]806 807 This is similar to :func:`substrings`, but emits items in a different808 order.809 """810 seq = list(iterable)811 slices = starmap(slice, combinations(range(len(seq) + 1), 2))812 return map(operator.getitem, repeat(seq), slices)813 814 815def polynomial_from_roots(roots):816 """Compute a polynomial's coefficients from its roots.817 818 >>> roots = [5, -4, 3] # (x - 5) * (x + 4) * (x - 3)819 >>> polynomial_from_roots(roots) # x^3 - 4 * x^2 - 17 * x + 60820 [1, -4, -17, 60]821 """822 factors = zip(repeat(1), map(operator.neg, roots))823 return list(reduce(convolve, factors, [1]))824 825 826def iter_index(iterable, value, start=0, stop=None):827 """Yield the index of each place in *iterable* that *value* occurs,828 beginning with index *start* and ending before index *stop*.829 830 See :func:`locate` for a more general means of finding the indexes831 associated with particular values.832 833 >>> list(iter_index('AABCADEAF', 'A'))834 [0, 1, 4, 7]835 >>> list(iter_index('AABCADEAF', 'A', 1)) # start index is inclusive836 [1, 4, 7]837 >>> list(iter_index('AABCADEAF', 'A', 1, 7)) # stop index is not inclusive838 [1, 4]839 """840 seq_index = getattr(iterable, 'index', None)841 if seq_index is None:842 # Slow path for general iterables843 it = islice(iterable, start, stop)844 for i, element in enumerate(it, start):845 if element is value or element == value:846 yield i847 else:848 # Fast path for sequences849 stop = len(iterable) if stop is None else stop850 i = start - 1851 try:852 while True:853 yield (i := seq_index(value, i + 1, stop))854 except ValueError:855 pass856 857 858def sieve(n):859 """Yield the primes less than n.860 861 >>> list(sieve(30))862 [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]863 """864 if n > 2:865 yield 2866 start = 3867 data = bytearray((0, 1)) * (n // 2)868 limit = math.isqrt(n) + 1869 for p in iter_index(data, 1, start, limit):870 yield from iter_index(data, 1, start, p * p)871 data[p * p : n : p + p] = bytes(len(range(p * p, n, p + p)))872 start = p * p873 yield from iter_index(data, 1, start)874 875 876def _batched(iterable, n, *, strict=False):877 """Batch data into tuples of length *n*. If the number of items in878 *iterable* is not divisible by *n*:879 * The last batch will be shorter if *strict* is ``False``.880 * :exc:`ValueError` will be raised if *strict* is ``True``.881 882 >>> list(batched('ABCDEFG', 3))883 [('A', 'B', 'C'), ('D', 'E', 'F'), ('G',)]884 885 On Python 3.13 and above, this is an alias for :func:`itertools.batched`.886 """887 if n < 1:888 raise ValueError('n must be at least one')889 it = iter(iterable)890 while batch := tuple(islice(it, n)):891 if strict and len(batch) != n:892 raise ValueError('batched(): incomplete batch')893 yield batch894 895 896if hexversion >= 0x30D00A2:897 from itertools import batched as itertools_batched898 899 def batched(iterable, n, *, strict=False):900 return itertools_batched(iterable, n, strict=strict)901 902else:903 batched = _batched904 905 batched.__doc__ = _batched.__doc__906 907 908def transpose(it):909 """Swap the rows and columns of the input matrix.910 911 >>> list(transpose([(1, 2, 3), (11, 22, 33)]))912 [(1, 11), (2, 22), (3, 33)]913 914 The caller should ensure that the dimensions of the input are compatible.915 If the input is empty, no output will be produced.916 """917 return _zip_strict(*it)918 919 920def reshape(matrix, cols):921 """Reshape the 2-D input *matrix* to have a column count given by *cols*.922 923 >>> matrix = [(0, 1), (2, 3), (4, 5)]924 >>> cols = 3925 >>> list(reshape(matrix, cols))926 [(0, 1, 2), (3, 4, 5)]927 """928 return batched(chain.from_iterable(matrix), cols)929 930 931def matmul(m1, m2):932 """Multiply two matrices.933 934 >>> list(matmul([(7, 5), (3, 5)], [(2, 5), (7, 9)]))935 [(49, 80), (41, 60)]936 937 The caller should ensure that the dimensions of the input matrices are938 compatible with each other.939 """940 n = len(m2[0])941 return batched(starmap(_sumprod, product(m1, transpose(m2))), n)942 943 944def factor(n):945 """Yield the prime factors of n.946 947 >>> list(factor(360))948 [2, 2, 2, 3, 3, 5]949 """950 for prime in sieve(math.isqrt(n) + 1):951 while not n % prime:952 yield prime953 n //= prime954 if n == 1:955 return956 if n > 1:957 yield n958 959 960def polynomial_eval(coefficients, x):961 """Evaluate a polynomial at a specific value.962 963 Example: evaluating x^3 - 4 * x^2 - 17 * x + 60 at x = 2.5:964 965 >>> coefficients = [1, -4, -17, 60]966 >>> x = 2.5967 >>> polynomial_eval(coefficients, x)968 8.125969 """970 n = len(coefficients)971 if n == 0:972 return x * 0 # coerce zero to the type of x973 powers = map(pow, repeat(x), reversed(range(n)))974 return _sumprod(coefficients, powers)975 976 977def sum_of_squares(it):978 """Return the sum of the squares of the input values.979 980 >>> sum_of_squares([10, 20, 30])981 1400982 """983 return _sumprod(*tee(it))984 985 986def polynomial_derivative(coefficients):987 """Compute the first derivative of a polynomial.988 989 Example: evaluating the derivative of x^3 - 4 * x^2 - 17 * x + 60990 991 >>> coefficients = [1, -4, -17, 60]992 >>> derivative_coefficients = polynomial_derivative(coefficients)993 >>> derivative_coefficients994 [3, -8, -17]995 """996 n = len(coefficients)997 powers = reversed(range(1, n))998 return list(map(operator.mul, coefficients, powers))999 1000 1001def totient(n):1002 """Return the count of natural numbers up to *n* that are coprime with *n*.1003 1004 >>> totient(9)1005 61006 >>> totient(12)1007 41008 """1009 for p in unique_justseen(factor(n)):1010 n = n // p * (p - 1)1011 1012 return n1013 