1 """Imported from the recipes section of the itertools documentation.
3 All functions taken from the recipes section of the itertools library docs
5 Some backward-compatible usability improvements have been made.
7 .. [1] http://docs.python.org/library/itertools.html#recipes
11 from collections import deque
12 from itertools import (
25 from random import randrange, sample, choice
46 'random_combination_with_replacement',
60 def take(n, iterable):
61 """Return first *n* items of the iterable as a list.
63 >>> take(3, range(10))
66 If there are fewer than *n* items in the iterable, all of them are
69 >>> take(10, range(3))
73 return list(islice(iterable, n))
76 def tabulate(function, start=0):
77 """Return an iterator over the results of ``func(start)``,
78 ``func(start + 1)``, ``func(start + 2)``...
80 *func* should be a function that accepts one integer argument.
82 If *start* is not specified it defaults to 0. It will be incremented each
83 time the iterator is advanced.
85 >>> square = lambda x: x ** 2
86 >>> iterator = tabulate(square, -3)
91 return map(function, count(start))
94 def tail(n, iterable):
95 """Return an iterator over the last *n* items of *iterable*.
97 >>> t = tail(3, 'ABCDEFG')
102 return iter(deque(iterable, maxlen=n))
105 def consume(iterator, n=None):
106 """Advance *iterable* by *n* steps. If *n* is ``None``, consume it
109 Efficiently exhausts an iterator without returning values. Defaults to
110 consuming the whole iterator, but an optional second argument may be
111 provided to limit consumption.
113 >>> i = (x for x in range(10))
121 Traceback (most recent call last):
122 File "<stdin>", line 1, in <module>
125 If the iterator has fewer items remaining than the provided limit, the
126 whole iterator will be consumed.
128 >>> i = (x for x in range(3))
131 Traceback (most recent call last):
132 File "<stdin>", line 1, in <module>
136 # Use functions that consume iterators at C speed.
138 # feed the entire iterator into a zero-length deque
139 deque(iterator, maxlen=0)
141 # advance to the empty slice starting at position n
142 next(islice(iterator, n, n), None)
145 def nth(iterable, n, default=None):
146 """Returns the nth item or a default value.
151 >>> nth(l, 20, "zebra")
155 return next(islice(iterable, n, None), default)
158 def all_equal(iterable):
160 Returns ``True`` if all the elements are equal to each other.
162 >>> all_equal('aaaa')
164 >>> all_equal('aaab')
168 g = groupby(iterable)
169 return next(g, True) and not next(g, False)
172 def quantify(iterable, pred=bool):
173 """Return the how many times the predicate is true.
175 >>> quantify([True, False, True])
179 return sum(map(pred, iterable))
182 def pad_none(iterable):
183 """Returns the sequence of elements and then returns ``None`` indefinitely.
185 >>> take(5, pad_none(range(3)))
186 [0, 1, 2, None, None]
188 Useful for emulating the behavior of the built-in :func:`map` function.
190 See also :func:`padded`.
193 return chain(iterable, repeat(None))
199 def ncycles(iterable, n):
200 """Returns the sequence elements *n* times
202 >>> list(ncycles(["a", "b"], 3))
203 ['a', 'b', 'a', 'b', 'a', 'b']
206 return chain.from_iterable(repeat(tuple(iterable), n))
209 def dotproduct(vec1, vec2):
210 """Returns the dot product of the two iterables.
212 >>> dotproduct([10, 10], [20, 20])
216 return sum(map(operator.mul, vec1, vec2))
219 def flatten(listOfLists):
220 """Return an iterator flattening one level of nesting in a list of lists.
222 >>> list(flatten([[0, 1], [2, 3]]))
225 See also :func:`collapse`, which can flatten multiple levels of nesting.
228 return chain.from_iterable(listOfLists)
231 def repeatfunc(func, times=None, *args):
232 """Call *func* with *args* repeatedly, returning an iterable over the
235 If *times* is specified, the iterable will terminate after that many
238 >>> from operator import add
241 >>> list(repeatfunc(add, times, *args))
244 If *times* is ``None`` the iterable will not terminate:
246 >>> from random import randrange
249 >>> take(6, repeatfunc(randrange, times, *args)) # doctest:+SKIP
254 return starmap(func, repeat(args))
255 return starmap(func, repeat(args, times))
258 def _pairwise(iterable):
259 """Returns an iterator of paired items, overlapping, from the original
261 >>> take(4, pairwise(count()))
262 [(0, 1), (1, 2), (2, 3), (3, 4)]
264 On Python 3.10 and above, this is an alias for :func:`itertools.pairwise`.
273 from itertools import pairwise as itertools_pairwise
278 def pairwise(iterable):
279 yield from itertools_pairwise(iterable)
281 pairwise.__doc__ = _pairwise.__doc__
284 def grouper(iterable, n, fillvalue=None):
285 """Collect data into fixed-length chunks or blocks.
287 >>> list(grouper('ABCDEFG', 3, 'x'))
288 [('A', 'B', 'C'), ('D', 'E', 'F'), ('G', 'x', 'x')]
291 if isinstance(iterable, int):
293 "grouper expects iterable as first parameter", DeprecationWarning
295 n, iterable = iterable, n
296 args = [iter(iterable)] * n
297 return zip_longest(fillvalue=fillvalue, *args)
300 def roundrobin(*iterables):
301 """Yields an item from each iterable, alternating between them.
303 >>> list(roundrobin('ABC', 'D', 'EF'))
304 ['A', 'D', 'E', 'B', 'F', 'C']
306 This function produces the same output as :func:`interleave_longest`, but
307 may perform better for some inputs (in particular when the number of
311 # Recipe credited to George Sakkis
312 pending = len(iterables)
313 nexts = cycle(iter(it).__next__ for it in iterables)
318 except StopIteration:
320 nexts = cycle(islice(nexts, pending))
323 def partition(pred, iterable):
325 Returns a 2-tuple of iterables derived from the input iterable.
326 The first yields the items that have ``pred(item) == False``.
327 The second yields the items that have ``pred(item) == True``.
329 >>> is_odd = lambda x: x % 2 != 0
330 >>> iterable = range(10)
331 >>> even_items, odd_items = partition(is_odd, iterable)
332 >>> list(even_items), list(odd_items)
333 ([0, 2, 4, 6, 8], [1, 3, 5, 7, 9])
335 If *pred* is None, :func:`bool` is used.
337 >>> iterable = [0, 1, False, True, '', ' ']
338 >>> false_items, true_items = partition(None, iterable)
339 >>> list(false_items), list(true_items)
340 ([0, False, ''], [1, True, ' '])
346 evaluations = ((pred(x), x) for x in iterable)
347 t1, t2 = tee(evaluations)
349 (x for (cond, x) in t1 if not cond),
350 (x for (cond, x) in t2 if cond),
354 def powerset(iterable):
355 """Yields all possible subsets of the iterable.
357 >>> list(powerset([1, 2, 3]))
358 [(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
360 :func:`powerset` will operate on iterables that aren't :class:`set`
361 instances, so repeated elements in the input will produce repeated elements
362 in the output. Use :func:`unique_everseen` on the input to avoid generating
366 >>> list(powerset(seq))
367 [(), (1,), (1,), (0,), (1, 1), (1, 0), (1, 0), (1, 1, 0)]
368 >>> from more_itertools import unique_everseen
369 >>> list(powerset(unique_everseen(seq)))
370 [(), (1,), (0,), (1, 0)]
374 return chain.from_iterable(combinations(s, r) for r in range(len(s) + 1))
377 def unique_everseen(iterable, key=None):
379 Yield unique elements, preserving order.
381 >>> list(unique_everseen('AAAABBBCCDAABBB'))
383 >>> list(unique_everseen('ABBCcAD', str.lower))
386 Sequences with a mix of hashable and unhashable items can be used.
387 The function will be slower (i.e., `O(n^2)`) for unhashable items.
389 Remember that ``list`` objects are unhashable - you can use the *key*
390 parameter to transform the list to a tuple (which is hashable) to
393 >>> iterable = ([1, 2], [2, 3], [1, 2])
394 >>> list(unique_everseen(iterable)) # Slow
396 >>> list(unique_everseen(iterable, key=tuple)) # Faster
399 Similary, you may want to convert unhashable ``set`` objects with
400 ``key=frozenset``. For ``dict`` objects,
401 ``key=lambda x: frozenset(x.items())`` can be used.
405 seenset_add = seenset.add
407 seenlist_add = seenlist.append
408 use_key = key is not None
410 for element in iterable:
411 k = key(element) if use_key else element
417 if k not in seenlist:
422 def unique_justseen(iterable, key=None):
423 """Yields elements in order, ignoring serial duplicates
425 >>> list(unique_justseen('AAAABBBCCDAABBB'))
426 ['A', 'B', 'C', 'D', 'A', 'B']
427 >>> list(unique_justseen('ABBCcAD', str.lower))
428 ['A', 'B', 'C', 'A', 'D']
431 return map(next, map(operator.itemgetter(1), groupby(iterable, key)))
434 def iter_except(func, exception, first=None):
435 """Yields results from a function repeatedly until an exception is raised.
437 Converts a call-until-exception interface to an iterator interface.
438 Like ``iter(func, sentinel)``, but uses an exception instead of a sentinel
442 >>> list(iter_except(l.pop, IndexError))
447 if first is not None:
455 def first_true(iterable, default=None, pred=None):
457 Returns the first true value in the iterable.
459 If no true value is found, returns *default*
461 If *pred* is not None, returns the first item for which
462 ``pred(item) == True`` .
464 >>> first_true(range(10))
466 >>> first_true(range(10), pred=lambda x: x > 5)
468 >>> first_true(range(10), default='missing', pred=lambda x: x > 9)
472 return next(filter(pred, iterable), default)
475 def random_product(*args, repeat=1):
476 """Draw an item at random from each of the input iterables.
478 >>> random_product('abc', range(4), 'XYZ') # doctest:+SKIP
481 If *repeat* is provided as a keyword argument, that many items will be
482 drawn from each iterable.
484 >>> random_product('abcd', range(4), repeat=2) # doctest:+SKIP
487 This equivalent to taking a random selection from
488 ``itertools.product(*args, **kwarg)``.
491 pools = [tuple(pool) for pool in args] * repeat
492 return tuple(choice(pool) for pool in pools)
495 def random_permutation(iterable, r=None):
496 """Return a random *r* length permutation of the elements in *iterable*.
498 If *r* is not specified or is ``None``, then *r* defaults to the length of
501 >>> random_permutation(range(5)) # doctest:+SKIP
504 This equivalent to taking a random selection from
505 ``itertools.permutations(iterable, r)``.
508 pool = tuple(iterable)
509 r = len(pool) if r is None else r
510 return tuple(sample(pool, r))
513 def random_combination(iterable, r):
514 """Return a random *r* length subsequence of the elements in *iterable*.
516 >>> random_combination(range(5), 3) # doctest:+SKIP
519 This equivalent to taking a random selection from
520 ``itertools.combinations(iterable, r)``.
523 pool = tuple(iterable)
525 indices = sorted(sample(range(n), r))
526 return tuple(pool[i] for i in indices)
529 def random_combination_with_replacement(iterable, r):
530 """Return a random *r* length subsequence of elements in *iterable*,
531 allowing individual elements to be repeated.
533 >>> random_combination_with_replacement(range(3), 5) # doctest:+SKIP
536 This equivalent to taking a random selection from
537 ``itertools.combinations_with_replacement(iterable, r)``.
540 pool = tuple(iterable)
542 indices = sorted(randrange(n) for i in range(r))
543 return tuple(pool[i] for i in indices)
546 def nth_combination(iterable, r, index):
547 """Equivalent to ``list(combinations(iterable, r))[index]``.
549 The subsequences of *iterable* that are of length *r* can be ordered
550 lexicographically. :func:`nth_combination` computes the subsequence at
551 sort position *index* directly, without computing the previous
554 >>> nth_combination(range(5), 3, 5)
557 ``ValueError`` will be raised If *r* is negative or greater than the length
559 ``IndexError`` will be raised if the given *index* is invalid.
561 pool = tuple(iterable)
563 if (r < 0) or (r > n):
568 for i in range(1, k + 1):
569 c = c * (n - k + i) // i
574 if (index < 0) or (index >= c):
579 c, n, r = c * r // n, n - 1, r - 1
582 c, n = c * (n - r) // n, n - 1
583 result.append(pool[-1 - n])
588 def prepend(value, iterator):
589 """Yield *value*, followed by the elements in *iterator*.
592 >>> iterator = ['1', '2', '3']
593 >>> list(prepend(value, iterator))
596 To prepend multiple values, see :func:`itertools.chain`
597 or :func:`value_chain`.
600 return chain([value], iterator)
603 def convolve(signal, kernel):
604 """Convolve the iterable *signal* with the iterable *kernel*.
606 >>> signal = (1, 2, 3, 4, 5)
607 >>> kernel = [3, 2, 1]
608 >>> list(convolve(signal, kernel))
609 [3, 8, 14, 20, 26, 14, 5]
611 Note: the input arguments are not interchangeable, as the *kernel*
612 is immediately consumed and stored.
615 kernel = tuple(kernel)[::-1]
617 window = deque([0], maxlen=n) * n
618 for x in chain(signal, repeat(0, n - 1)):
620 yield sum(map(operator.mul, kernel, window))