Team Ai
Datasetpublic

codekingpro/portable-devtools

sourceHugging Faceupdated 5mo agoView on Hugging Face
1likes14kdownloads
recipes.py1013 linesDownload Raw Back to more_itertools
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 
codekingpro/portable-devtools · Team Ai