recipes.py 29 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070
  1. """Imported from the recipes section of the itertools documentation.
  2. All functions taken from the recipes section of the itertools library docs
  3. [1]_.
  4. Some backward-compatible usability improvements have been made.
  5. .. [1] http://docs.python.org/library/itertools.html#recipes
  6. """
  7. import math
  8. import operator
  9. from collections import deque
  10. from collections.abc import Sized
  11. from functools import partial, reduce
  12. from itertools import (
  13. chain,
  14. combinations,
  15. compress,
  16. count,
  17. cycle,
  18. groupby,
  19. islice,
  20. product,
  21. repeat,
  22. starmap,
  23. tee,
  24. zip_longest,
  25. )
  26. from random import randrange, sample, choice
  27. from sys import hexversion
  28. __all__ = [
  29. 'all_equal',
  30. 'batched',
  31. 'before_and_after',
  32. 'consume',
  33. 'convolve',
  34. 'dotproduct',
  35. 'first_true',
  36. 'factor',
  37. 'flatten',
  38. 'grouper',
  39. 'iter_except',
  40. 'iter_index',
  41. 'matmul',
  42. 'ncycles',
  43. 'nth',
  44. 'nth_combination',
  45. 'padnone',
  46. 'pad_none',
  47. 'pairwise',
  48. 'partition',
  49. 'polynomial_eval',
  50. 'polynomial_from_roots',
  51. 'polynomial_derivative',
  52. 'powerset',
  53. 'prepend',
  54. 'quantify',
  55. 'reshape',
  56. 'random_combination_with_replacement',
  57. 'random_combination',
  58. 'random_permutation',
  59. 'random_product',
  60. 'repeatfunc',
  61. 'roundrobin',
  62. 'sieve',
  63. 'sliding_window',
  64. 'subslices',
  65. 'sum_of_squares',
  66. 'tabulate',
  67. 'tail',
  68. 'take',
  69. 'totient',
  70. 'transpose',
  71. 'triplewise',
  72. 'unique',
  73. 'unique_everseen',
  74. 'unique_justseen',
  75. ]
  76. _marker = object()
  77. # zip with strict is available for Python 3.10+
  78. try:
  79. zip(strict=True)
  80. except TypeError:
  81. _zip_strict = zip
  82. else:
  83. _zip_strict = partial(zip, strict=True)
  84. # math.sumprod is available for Python 3.12+
  85. _sumprod = getattr(math, 'sumprod', lambda x, y: dotproduct(x, y))
  86. def take(n, iterable):
  87. """Return first *n* items of the iterable as a list.
  88. >>> take(3, range(10))
  89. [0, 1, 2]
  90. If there are fewer than *n* items in the iterable, all of them are
  91. returned.
  92. >>> take(10, range(3))
  93. [0, 1, 2]
  94. """
  95. return list(islice(iterable, n))
  96. def tabulate(function, start=0):
  97. """Return an iterator over the results of ``func(start)``,
  98. ``func(start + 1)``, ``func(start + 2)``...
  99. *func* should be a function that accepts one integer argument.
  100. If *start* is not specified it defaults to 0. It will be incremented each
  101. time the iterator is advanced.
  102. >>> square = lambda x: x ** 2
  103. >>> iterator = tabulate(square, -3)
  104. >>> take(4, iterator)
  105. [9, 4, 1, 0]
  106. """
  107. return map(function, count(start))
  108. def tail(n, iterable):
  109. """Return an iterator over the last *n* items of *iterable*.
  110. >>> t = tail(3, 'ABCDEFG')
  111. >>> list(t)
  112. ['E', 'F', 'G']
  113. """
  114. # If the given iterable has a length, then we can use islice to get its
  115. # final elements. Note that if the iterable is not actually Iterable,
  116. # either islice or deque will throw a TypeError. This is why we don't
  117. # check if it is Iterable.
  118. if isinstance(iterable, Sized):
  119. yield from islice(iterable, max(0, len(iterable) - n), None)
  120. else:
  121. yield from iter(deque(iterable, maxlen=n))
  122. def consume(iterator, n=None):
  123. """Advance *iterable* by *n* steps. If *n* is ``None``, consume it
  124. entirely.
  125. Efficiently exhausts an iterator without returning values. Defaults to
  126. consuming the whole iterator, but an optional second argument may be
  127. provided to limit consumption.
  128. >>> i = (x for x in range(10))
  129. >>> next(i)
  130. 0
  131. >>> consume(i, 3)
  132. >>> next(i)
  133. 4
  134. >>> consume(i)
  135. >>> next(i)
  136. Traceback (most recent call last):
  137. File "<stdin>", line 1, in <module>
  138. StopIteration
  139. If the iterator has fewer items remaining than the provided limit, the
  140. whole iterator will be consumed.
  141. >>> i = (x for x in range(3))
  142. >>> consume(i, 5)
  143. >>> next(i)
  144. Traceback (most recent call last):
  145. File "<stdin>", line 1, in <module>
  146. StopIteration
  147. """
  148. # Use functions that consume iterators at C speed.
  149. if n is None:
  150. # feed the entire iterator into a zero-length deque
  151. deque(iterator, maxlen=0)
  152. else:
  153. # advance to the empty slice starting at position n
  154. next(islice(iterator, n, n), None)
  155. def nth(iterable, n, default=None):
  156. """Returns the nth item or a default value.
  157. >>> l = range(10)
  158. >>> nth(l, 3)
  159. 3
  160. >>> nth(l, 20, "zebra")
  161. 'zebra'
  162. """
  163. return next(islice(iterable, n, None), default)
  164. def all_equal(iterable, key=None):
  165. """
  166. Returns ``True`` if all the elements are equal to each other.
  167. >>> all_equal('aaaa')
  168. True
  169. >>> all_equal('aaab')
  170. False
  171. A function that accepts a single argument and returns a transformed version
  172. of each input item can be specified with *key*:
  173. >>> all_equal('AaaA', key=str.casefold)
  174. True
  175. >>> all_equal([1, 2, 3], key=lambda x: x < 10)
  176. True
  177. """
  178. return len(list(islice(groupby(iterable, key), 2))) <= 1
  179. def quantify(iterable, pred=bool):
  180. """Return the how many times the predicate is true.
  181. >>> quantify([True, False, True])
  182. 2
  183. """
  184. return sum(map(pred, iterable))
  185. def pad_none(iterable):
  186. """Returns the sequence of elements and then returns ``None`` indefinitely.
  187. >>> take(5, pad_none(range(3)))
  188. [0, 1, 2, None, None]
  189. Useful for emulating the behavior of the built-in :func:`map` function.
  190. See also :func:`padded`.
  191. """
  192. return chain(iterable, repeat(None))
  193. padnone = pad_none
  194. def ncycles(iterable, n):
  195. """Returns the sequence elements *n* times
  196. >>> list(ncycles(["a", "b"], 3))
  197. ['a', 'b', 'a', 'b', 'a', 'b']
  198. """
  199. return chain.from_iterable(repeat(tuple(iterable), n))
  200. def dotproduct(vec1, vec2):
  201. """Returns the dot product of the two iterables.
  202. >>> dotproduct([10, 10], [20, 20])
  203. 400
  204. """
  205. return sum(map(operator.mul, vec1, vec2))
  206. def flatten(listOfLists):
  207. """Return an iterator flattening one level of nesting in a list of lists.
  208. >>> list(flatten([[0, 1], [2, 3]]))
  209. [0, 1, 2, 3]
  210. See also :func:`collapse`, which can flatten multiple levels of nesting.
  211. """
  212. return chain.from_iterable(listOfLists)
  213. def repeatfunc(func, times=None, *args):
  214. """Call *func* with *args* repeatedly, returning an iterable over the
  215. results.
  216. If *times* is specified, the iterable will terminate after that many
  217. repetitions:
  218. >>> from operator import add
  219. >>> times = 4
  220. >>> args = 3, 5
  221. >>> list(repeatfunc(add, times, *args))
  222. [8, 8, 8, 8]
  223. If *times* is ``None`` the iterable will not terminate:
  224. >>> from random import randrange
  225. >>> times = None
  226. >>> args = 1, 11
  227. >>> take(6, repeatfunc(randrange, times, *args)) # doctest:+SKIP
  228. [2, 4, 8, 1, 8, 4]
  229. """
  230. if times is None:
  231. return starmap(func, repeat(args))
  232. return starmap(func, repeat(args, times))
  233. def _pairwise(iterable):
  234. """Returns an iterator of paired items, overlapping, from the original
  235. >>> take(4, pairwise(count()))
  236. [(0, 1), (1, 2), (2, 3), (3, 4)]
  237. On Python 3.10 and above, this is an alias for :func:`itertools.pairwise`.
  238. """
  239. a, b = tee(iterable)
  240. next(b, None)
  241. return zip(a, b)
  242. try:
  243. from itertools import pairwise as itertools_pairwise
  244. except ImportError:
  245. pairwise = _pairwise
  246. else:
  247. def pairwise(iterable):
  248. return itertools_pairwise(iterable)
  249. pairwise.__doc__ = _pairwise.__doc__
  250. class UnequalIterablesError(ValueError):
  251. def __init__(self, details=None):
  252. msg = 'Iterables have different lengths'
  253. if details is not None:
  254. msg += (': index 0 has length {}; index {} has length {}').format(
  255. *details
  256. )
  257. super().__init__(msg)
  258. def _zip_equal_generator(iterables):
  259. for combo in zip_longest(*iterables, fillvalue=_marker):
  260. for val in combo:
  261. if val is _marker:
  262. raise UnequalIterablesError()
  263. yield combo
  264. def _zip_equal(*iterables):
  265. # Check whether the iterables are all the same size.
  266. try:
  267. first_size = len(iterables[0])
  268. for i, it in enumerate(iterables[1:], 1):
  269. size = len(it)
  270. if size != first_size:
  271. raise UnequalIterablesError(details=(first_size, i, size))
  272. # All sizes are equal, we can use the built-in zip.
  273. return zip(*iterables)
  274. # If any one of the iterables didn't have a length, start reading
  275. # them until one runs out.
  276. except TypeError:
  277. return _zip_equal_generator(iterables)
  278. def grouper(iterable, n, incomplete='fill', fillvalue=None):
  279. """Group elements from *iterable* into fixed-length groups of length *n*.
  280. >>> list(grouper('ABCDEF', 3))
  281. [('A', 'B', 'C'), ('D', 'E', 'F')]
  282. The keyword arguments *incomplete* and *fillvalue* control what happens for
  283. iterables whose length is not a multiple of *n*.
  284. When *incomplete* is `'fill'`, the last group will contain instances of
  285. *fillvalue*.
  286. >>> list(grouper('ABCDEFG', 3, incomplete='fill', fillvalue='x'))
  287. [('A', 'B', 'C'), ('D', 'E', 'F'), ('G', 'x', 'x')]
  288. When *incomplete* is `'ignore'`, the last group will not be emitted.
  289. >>> list(grouper('ABCDEFG', 3, incomplete='ignore', fillvalue='x'))
  290. [('A', 'B', 'C'), ('D', 'E', 'F')]
  291. When *incomplete* is `'strict'`, a subclass of `ValueError` will be raised.
  292. >>> it = grouper('ABCDEFG', 3, incomplete='strict')
  293. >>> list(it) # doctest: +IGNORE_EXCEPTION_DETAIL
  294. Traceback (most recent call last):
  295. ...
  296. UnequalIterablesError
  297. """
  298. args = [iter(iterable)] * n
  299. if incomplete == 'fill':
  300. return zip_longest(*args, fillvalue=fillvalue)
  301. if incomplete == 'strict':
  302. return _zip_equal(*args)
  303. if incomplete == 'ignore':
  304. return zip(*args)
  305. else:
  306. raise ValueError('Expected fill, strict, or ignore')
  307. def roundrobin(*iterables):
  308. """Yields an item from each iterable, alternating between them.
  309. >>> list(roundrobin('ABC', 'D', 'EF'))
  310. ['A', 'D', 'E', 'B', 'F', 'C']
  311. This function produces the same output as :func:`interleave_longest`, but
  312. may perform better for some inputs (in particular when the number of
  313. iterables is small).
  314. """
  315. # Algorithm credited to George Sakkis
  316. iterators = map(iter, iterables)
  317. for num_active in range(len(iterables), 0, -1):
  318. iterators = cycle(islice(iterators, num_active))
  319. yield from map(next, iterators)
  320. def partition(pred, iterable):
  321. """
  322. Returns a 2-tuple of iterables derived from the input iterable.
  323. The first yields the items that have ``pred(item) == False``.
  324. The second yields the items that have ``pred(item) == True``.
  325. >>> is_odd = lambda x: x % 2 != 0
  326. >>> iterable = range(10)
  327. >>> even_items, odd_items = partition(is_odd, iterable)
  328. >>> list(even_items), list(odd_items)
  329. ([0, 2, 4, 6, 8], [1, 3, 5, 7, 9])
  330. If *pred* is None, :func:`bool` is used.
  331. >>> iterable = [0, 1, False, True, '', ' ']
  332. >>> false_items, true_items = partition(None, iterable)
  333. >>> list(false_items), list(true_items)
  334. ([0, False, ''], [1, True, ' '])
  335. """
  336. if pred is None:
  337. pred = bool
  338. t1, t2, p = tee(iterable, 3)
  339. p1, p2 = tee(map(pred, p))
  340. return (compress(t1, map(operator.not_, p1)), compress(t2, p2))
  341. def powerset(iterable):
  342. """Yields all possible subsets of the iterable.
  343. >>> list(powerset([1, 2, 3]))
  344. [(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
  345. :func:`powerset` will operate on iterables that aren't :class:`set`
  346. instances, so repeated elements in the input will produce repeated elements
  347. in the output.
  348. >>> seq = [1, 1, 0]
  349. >>> list(powerset(seq))
  350. [(), (1,), (1,), (0,), (1, 1), (1, 0), (1, 0), (1, 1, 0)]
  351. For a variant that efficiently yields actual :class:`set` instances, see
  352. :func:`powerset_of_sets`.
  353. """
  354. s = list(iterable)
  355. return chain.from_iterable(combinations(s, r) for r in range(len(s) + 1))
  356. def unique_everseen(iterable, key=None):
  357. """
  358. Yield unique elements, preserving order.
  359. >>> list(unique_everseen('AAAABBBCCDAABBB'))
  360. ['A', 'B', 'C', 'D']
  361. >>> list(unique_everseen('ABBCcAD', str.lower))
  362. ['A', 'B', 'C', 'D']
  363. Sequences with a mix of hashable and unhashable items can be used.
  364. The function will be slower (i.e., `O(n^2)`) for unhashable items.
  365. Remember that ``list`` objects are unhashable - you can use the *key*
  366. parameter to transform the list to a tuple (which is hashable) to
  367. avoid a slowdown.
  368. >>> iterable = ([1, 2], [2, 3], [1, 2])
  369. >>> list(unique_everseen(iterable)) # Slow
  370. [[1, 2], [2, 3]]
  371. >>> list(unique_everseen(iterable, key=tuple)) # Faster
  372. [[1, 2], [2, 3]]
  373. Similarly, you may want to convert unhashable ``set`` objects with
  374. ``key=frozenset``. For ``dict`` objects,
  375. ``key=lambda x: frozenset(x.items())`` can be used.
  376. """
  377. seenset = set()
  378. seenset_add = seenset.add
  379. seenlist = []
  380. seenlist_add = seenlist.append
  381. use_key = key is not None
  382. for element in iterable:
  383. k = key(element) if use_key else element
  384. try:
  385. if k not in seenset:
  386. seenset_add(k)
  387. yield element
  388. except TypeError:
  389. if k not in seenlist:
  390. seenlist_add(k)
  391. yield element
  392. def unique_justseen(iterable, key=None):
  393. """Yields elements in order, ignoring serial duplicates
  394. >>> list(unique_justseen('AAAABBBCCDAABBB'))
  395. ['A', 'B', 'C', 'D', 'A', 'B']
  396. >>> list(unique_justseen('ABBCcAD', str.lower))
  397. ['A', 'B', 'C', 'A', 'D']
  398. """
  399. if key is None:
  400. return map(operator.itemgetter(0), groupby(iterable))
  401. return map(next, map(operator.itemgetter(1), groupby(iterable, key)))
  402. def unique(iterable, key=None, reverse=False):
  403. """Yields unique elements in sorted order.
  404. >>> list(unique([[1, 2], [3, 4], [1, 2]]))
  405. [[1, 2], [3, 4]]
  406. *key* and *reverse* are passed to :func:`sorted`.
  407. >>> list(unique('ABBcCAD', str.casefold))
  408. ['A', 'B', 'c', 'D']
  409. >>> list(unique('ABBcCAD', str.casefold, reverse=True))
  410. ['D', 'c', 'B', 'A']
  411. The elements in *iterable* need not be hashable, but they must be
  412. comparable for sorting to work.
  413. """
  414. return unique_justseen(sorted(iterable, key=key, reverse=reverse), key=key)
  415. def iter_except(func, exception, first=None):
  416. """Yields results from a function repeatedly until an exception is raised.
  417. Converts a call-until-exception interface to an iterator interface.
  418. Like ``iter(func, sentinel)``, but uses an exception instead of a sentinel
  419. to end the loop.
  420. >>> l = [0, 1, 2]
  421. >>> list(iter_except(l.pop, IndexError))
  422. [2, 1, 0]
  423. Multiple exceptions can be specified as a stopping condition:
  424. >>> l = [1, 2, 3, '...', 4, 5, 6]
  425. >>> list(iter_except(lambda: 1 + l.pop(), (IndexError, TypeError)))
  426. [7, 6, 5]
  427. >>> list(iter_except(lambda: 1 + l.pop(), (IndexError, TypeError)))
  428. [4, 3, 2]
  429. >>> list(iter_except(lambda: 1 + l.pop(), (IndexError, TypeError)))
  430. []
  431. """
  432. try:
  433. if first is not None:
  434. yield first()
  435. while 1:
  436. yield func()
  437. except exception:
  438. pass
  439. def first_true(iterable, default=None, pred=None):
  440. """
  441. Returns the first true value in the iterable.
  442. If no true value is found, returns *default*
  443. If *pred* is not None, returns the first item for which
  444. ``pred(item) == True`` .
  445. >>> first_true(range(10))
  446. 1
  447. >>> first_true(range(10), pred=lambda x: x > 5)
  448. 6
  449. >>> first_true(range(10), default='missing', pred=lambda x: x > 9)
  450. 'missing'
  451. """
  452. return next(filter(pred, iterable), default)
  453. def random_product(*args, repeat=1):
  454. """Draw an item at random from each of the input iterables.
  455. >>> random_product('abc', range(4), 'XYZ') # doctest:+SKIP
  456. ('c', 3, 'Z')
  457. If *repeat* is provided as a keyword argument, that many items will be
  458. drawn from each iterable.
  459. >>> random_product('abcd', range(4), repeat=2) # doctest:+SKIP
  460. ('a', 2, 'd', 3)
  461. This equivalent to taking a random selection from
  462. ``itertools.product(*args, **kwarg)``.
  463. """
  464. pools = [tuple(pool) for pool in args] * repeat
  465. return tuple(choice(pool) for pool in pools)
  466. def random_permutation(iterable, r=None):
  467. """Return a random *r* length permutation of the elements in *iterable*.
  468. If *r* is not specified or is ``None``, then *r* defaults to the length of
  469. *iterable*.
  470. >>> random_permutation(range(5)) # doctest:+SKIP
  471. (3, 4, 0, 1, 2)
  472. This equivalent to taking a random selection from
  473. ``itertools.permutations(iterable, r)``.
  474. """
  475. pool = tuple(iterable)
  476. r = len(pool) if r is None else r
  477. return tuple(sample(pool, r))
  478. def random_combination(iterable, r):
  479. """Return a random *r* length subsequence of the elements in *iterable*.
  480. >>> random_combination(range(5), 3) # doctest:+SKIP
  481. (2, 3, 4)
  482. This equivalent to taking a random selection from
  483. ``itertools.combinations(iterable, r)``.
  484. """
  485. pool = tuple(iterable)
  486. n = len(pool)
  487. indices = sorted(sample(range(n), r))
  488. return tuple(pool[i] for i in indices)
  489. def random_combination_with_replacement(iterable, r):
  490. """Return a random *r* length subsequence of elements in *iterable*,
  491. allowing individual elements to be repeated.
  492. >>> random_combination_with_replacement(range(3), 5) # doctest:+SKIP
  493. (0, 0, 1, 2, 2)
  494. This equivalent to taking a random selection from
  495. ``itertools.combinations_with_replacement(iterable, r)``.
  496. """
  497. pool = tuple(iterable)
  498. n = len(pool)
  499. indices = sorted(randrange(n) for i in range(r))
  500. return tuple(pool[i] for i in indices)
  501. def nth_combination(iterable, r, index):
  502. """Equivalent to ``list(combinations(iterable, r))[index]``.
  503. The subsequences of *iterable* that are of length *r* can be ordered
  504. lexicographically. :func:`nth_combination` computes the subsequence at
  505. sort position *index* directly, without computing the previous
  506. subsequences.
  507. >>> nth_combination(range(5), 3, 5)
  508. (0, 3, 4)
  509. ``ValueError`` will be raised If *r* is negative or greater than the length
  510. of *iterable*.
  511. ``IndexError`` will be raised if the given *index* is invalid.
  512. """
  513. pool = tuple(iterable)
  514. n = len(pool)
  515. if (r < 0) or (r > n):
  516. raise ValueError
  517. c = 1
  518. k = min(r, n - r)
  519. for i in range(1, k + 1):
  520. c = c * (n - k + i) // i
  521. if index < 0:
  522. index += c
  523. if (index < 0) or (index >= c):
  524. raise IndexError
  525. result = []
  526. while r:
  527. c, n, r = c * r // n, n - 1, r - 1
  528. while index >= c:
  529. index -= c
  530. c, n = c * (n - r) // n, n - 1
  531. result.append(pool[-1 - n])
  532. return tuple(result)
  533. def prepend(value, iterator):
  534. """Yield *value*, followed by the elements in *iterator*.
  535. >>> value = '0'
  536. >>> iterator = ['1', '2', '3']
  537. >>> list(prepend(value, iterator))
  538. ['0', '1', '2', '3']
  539. To prepend multiple values, see :func:`itertools.chain`
  540. or :func:`value_chain`.
  541. """
  542. return chain([value], iterator)
  543. def convolve(signal, kernel):
  544. """Convolve the iterable *signal* with the iterable *kernel*.
  545. >>> signal = (1, 2, 3, 4, 5)
  546. >>> kernel = [3, 2, 1]
  547. >>> list(convolve(signal, kernel))
  548. [3, 8, 14, 20, 26, 14, 5]
  549. Note: the input arguments are not interchangeable, as the *kernel*
  550. is immediately consumed and stored.
  551. """
  552. # This implementation intentionally doesn't match the one in the itertools
  553. # documentation.
  554. kernel = tuple(kernel)[::-1]
  555. n = len(kernel)
  556. window = deque([0], maxlen=n) * n
  557. for x in chain(signal, repeat(0, n - 1)):
  558. window.append(x)
  559. yield _sumprod(kernel, window)
  560. def before_and_after(predicate, it):
  561. """A variant of :func:`takewhile` that allows complete access to the
  562. remainder of the iterator.
  563. >>> it = iter('ABCdEfGhI')
  564. >>> all_upper, remainder = before_and_after(str.isupper, it)
  565. >>> ''.join(all_upper)
  566. 'ABC'
  567. >>> ''.join(remainder) # takewhile() would lose the 'd'
  568. 'dEfGhI'
  569. Note that the first iterator must be fully consumed before the second
  570. iterator can generate valid results.
  571. """
  572. it = iter(it)
  573. transition = []
  574. def true_iterator():
  575. for elem in it:
  576. if predicate(elem):
  577. yield elem
  578. else:
  579. transition.append(elem)
  580. return
  581. # Note: this is different from itertools recipes to allow nesting
  582. # before_and_after remainders into before_and_after again. See tests
  583. # for an example.
  584. remainder_iterator = chain(transition, it)
  585. return true_iterator(), remainder_iterator
  586. def triplewise(iterable):
  587. """Return overlapping triplets from *iterable*.
  588. >>> list(triplewise('ABCDE'))
  589. [('A', 'B', 'C'), ('B', 'C', 'D'), ('C', 'D', 'E')]
  590. """
  591. # This deviates from the itertools documentation reciple - see
  592. # https://github.com/more-itertools/more-itertools/issues/889
  593. t1, t2, t3 = tee(iterable, 3)
  594. next(t3, None)
  595. next(t3, None)
  596. next(t2, None)
  597. return zip(t1, t2, t3)
  598. def _sliding_window_islice(iterable, n):
  599. # Fast path for small, non-zero values of n.
  600. iterators = tee(iterable, n)
  601. for i, iterator in enumerate(iterators):
  602. next(islice(iterator, i, i), None)
  603. return zip(*iterators)
  604. def _sliding_window_deque(iterable, n):
  605. # Normal path for other values of n.
  606. it = iter(iterable)
  607. window = deque(islice(it, n - 1), maxlen=n)
  608. for x in it:
  609. window.append(x)
  610. yield tuple(window)
  611. def sliding_window(iterable, n):
  612. """Return a sliding window of width *n* over *iterable*.
  613. >>> list(sliding_window(range(6), 4))
  614. [(0, 1, 2, 3), (1, 2, 3, 4), (2, 3, 4, 5)]
  615. If *iterable* has fewer than *n* items, then nothing is yielded:
  616. >>> list(sliding_window(range(3), 4))
  617. []
  618. For a variant with more features, see :func:`windowed`.
  619. """
  620. if n > 20:
  621. return _sliding_window_deque(iterable, n)
  622. elif n > 2:
  623. return _sliding_window_islice(iterable, n)
  624. elif n == 2:
  625. return pairwise(iterable)
  626. elif n == 1:
  627. return zip(iterable)
  628. else:
  629. raise ValueError(f'n should be at least one, not {n}')
  630. def subslices(iterable):
  631. """Return all contiguous non-empty subslices of *iterable*.
  632. >>> list(subslices('ABC'))
  633. [['A'], ['A', 'B'], ['A', 'B', 'C'], ['B'], ['B', 'C'], ['C']]
  634. This is similar to :func:`substrings`, but emits items in a different
  635. order.
  636. """
  637. seq = list(iterable)
  638. slices = starmap(slice, combinations(range(len(seq) + 1), 2))
  639. return map(operator.getitem, repeat(seq), slices)
  640. def polynomial_from_roots(roots):
  641. """Compute a polynomial's coefficients from its roots.
  642. >>> roots = [5, -4, 3] # (x - 5) * (x + 4) * (x - 3)
  643. >>> polynomial_from_roots(roots) # x^3 - 4 * x^2 - 17 * x + 60
  644. [1, -4, -17, 60]
  645. """
  646. factors = zip(repeat(1), map(operator.neg, roots))
  647. return list(reduce(convolve, factors, [1]))
  648. def iter_index(iterable, value, start=0, stop=None):
  649. """Yield the index of each place in *iterable* that *value* occurs,
  650. beginning with index *start* and ending before index *stop*.
  651. >>> list(iter_index('AABCADEAF', 'A'))
  652. [0, 1, 4, 7]
  653. >>> list(iter_index('AABCADEAF', 'A', 1)) # start index is inclusive
  654. [1, 4, 7]
  655. >>> list(iter_index('AABCADEAF', 'A', 1, 7)) # stop index is not inclusive
  656. [1, 4]
  657. The behavior for non-scalar *values* matches the built-in Python types.
  658. >>> list(iter_index('ABCDABCD', 'AB'))
  659. [0, 4]
  660. >>> list(iter_index([0, 1, 2, 3, 0, 1, 2, 3], [0, 1]))
  661. []
  662. >>> list(iter_index([[0, 1], [2, 3], [0, 1], [2, 3]], [0, 1]))
  663. [0, 2]
  664. See :func:`locate` for a more general means of finding the indexes
  665. associated with particular values.
  666. """
  667. seq_index = getattr(iterable, 'index', None)
  668. if seq_index is None:
  669. # Slow path for general iterables
  670. it = islice(iterable, start, stop)
  671. for i, element in enumerate(it, start):
  672. if element is value or element == value:
  673. yield i
  674. else:
  675. # Fast path for sequences
  676. stop = len(iterable) if stop is None else stop
  677. i = start - 1
  678. try:
  679. while True:
  680. yield (i := seq_index(value, i + 1, stop))
  681. except ValueError:
  682. pass
  683. def sieve(n):
  684. """Yield the primes less than n.
  685. >>> list(sieve(30))
  686. [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]
  687. """
  688. if n > 2:
  689. yield 2
  690. start = 3
  691. data = bytearray((0, 1)) * (n // 2)
  692. limit = math.isqrt(n) + 1
  693. for p in iter_index(data, 1, start, limit):
  694. yield from iter_index(data, 1, start, p * p)
  695. data[p * p : n : p + p] = bytes(len(range(p * p, n, p + p)))
  696. start = p * p
  697. yield from iter_index(data, 1, start)
  698. def _batched(iterable, n, *, strict=False):
  699. """Batch data into tuples of length *n*. If the number of items in
  700. *iterable* is not divisible by *n*:
  701. * The last batch will be shorter if *strict* is ``False``.
  702. * :exc:`ValueError` will be raised if *strict* is ``True``.
  703. >>> list(batched('ABCDEFG', 3))
  704. [('A', 'B', 'C'), ('D', 'E', 'F'), ('G',)]
  705. On Python 3.13 and above, this is an alias for :func:`itertools.batched`.
  706. """
  707. if n < 1:
  708. raise ValueError('n must be at least one')
  709. it = iter(iterable)
  710. while batch := tuple(islice(it, n)):
  711. if strict and len(batch) != n:
  712. raise ValueError('batched(): incomplete batch')
  713. yield batch
  714. if hexversion >= 0x30D00A2:
  715. from itertools import batched as itertools_batched
  716. def batched(iterable, n, *, strict=False):
  717. return itertools_batched(iterable, n, strict=strict)
  718. else:
  719. batched = _batched
  720. batched.__doc__ = _batched.__doc__
  721. def transpose(it):
  722. """Swap the rows and columns of the input matrix.
  723. >>> list(transpose([(1, 2, 3), (11, 22, 33)]))
  724. [(1, 11), (2, 22), (3, 33)]
  725. The caller should ensure that the dimensions of the input are compatible.
  726. If the input is empty, no output will be produced.
  727. """
  728. return _zip_strict(*it)
  729. def reshape(matrix, cols):
  730. """Reshape the 2-D input *matrix* to have a column count given by *cols*.
  731. >>> matrix = [(0, 1), (2, 3), (4, 5)]
  732. >>> cols = 3
  733. >>> list(reshape(matrix, cols))
  734. [(0, 1, 2), (3, 4, 5)]
  735. """
  736. return batched(chain.from_iterable(matrix), cols)
  737. def matmul(m1, m2):
  738. """Multiply two matrices.
  739. >>> list(matmul([(7, 5), (3, 5)], [(2, 5), (7, 9)]))
  740. [(49, 80), (41, 60)]
  741. The caller should ensure that the dimensions of the input matrices are
  742. compatible with each other.
  743. """
  744. n = len(m2[0])
  745. return batched(starmap(_sumprod, product(m1, transpose(m2))), n)
  746. def factor(n):
  747. """Yield the prime factors of n.
  748. >>> list(factor(360))
  749. [2, 2, 2, 3, 3, 5]
  750. """
  751. for prime in sieve(math.isqrt(n) + 1):
  752. while not n % prime:
  753. yield prime
  754. n //= prime
  755. if n == 1:
  756. return
  757. if n > 1:
  758. yield n
  759. def polynomial_eval(coefficients, x):
  760. """Evaluate a polynomial at a specific value.
  761. Example: evaluating x^3 - 4 * x^2 - 17 * x + 60 at x = 2.5:
  762. >>> coefficients = [1, -4, -17, 60]
  763. >>> x = 2.5
  764. >>> polynomial_eval(coefficients, x)
  765. 8.125
  766. """
  767. n = len(coefficients)
  768. if n == 0:
  769. return x * 0 # coerce zero to the type of x
  770. powers = map(pow, repeat(x), reversed(range(n)))
  771. return _sumprod(coefficients, powers)
  772. def sum_of_squares(it):
  773. """Return the sum of the squares of the input values.
  774. >>> sum_of_squares([10, 20, 30])
  775. 1400
  776. """
  777. return _sumprod(*tee(it))
  778. def polynomial_derivative(coefficients):
  779. """Compute the first derivative of a polynomial.
  780. Example: evaluating the derivative of x^3 - 4 * x^2 - 17 * x + 60
  781. >>> coefficients = [1, -4, -17, 60]
  782. >>> derivative_coefficients = polynomial_derivative(coefficients)
  783. >>> derivative_coefficients
  784. [3, -8, -17]
  785. """
  786. n = len(coefficients)
  787. powers = reversed(range(1, n))
  788. return list(map(operator.mul, coefficients, powers))
  789. def totient(n):
  790. """Return the count of natural numbers up to *n* that are coprime with *n*.
  791. >>> totient(9)
  792. 6
  793. >>> totient(12)
  794. 4
  795. """
  796. for prime in set(factor(n)):
  797. n -= n // prime
  798. return n