blob: 959424ff914390e8151089da5dfaa24a8aa9433e [file] [log] [blame]
Georg Brandl116aa622007-08-15 14:28:22 +00001:mod:`itertools` --- Functions creating iterators for efficient looping
2=======================================================================
3
4.. module:: itertools
Raymond Hettinger30c73622010-12-01 23:45:20 +00005 :synopsis: Functions creating iterators for efficient looping.
Terry Jan Reedyfa089b92016-06-11 15:02:54 -04006
Georg Brandl116aa622007-08-15 14:28:22 +00007.. moduleauthor:: Raymond Hettinger <python@rcn.com>
8.. sectionauthor:: Raymond Hettinger <python@rcn.com>
9
Christian Heimesfe337bf2008-03-23 21:54:12 +000010.. testsetup::
11
Raymond Hettinger30c73622010-12-01 23:45:20 +000012 from itertools import *
Christian Heimesfe337bf2008-03-23 21:54:12 +000013
Terry Jan Reedyfa089b92016-06-11 15:02:54 -040014--------------
Christian Heimesfe337bf2008-03-23 21:54:12 +000015
Raymond Hettingerf76b9202009-02-17 20:00:59 +000016This module implements a number of :term:`iterator` building blocks inspired
17by constructs from APL, Haskell, and SML. Each has been recast in a form
18suitable for Python.
Georg Brandl116aa622007-08-15 14:28:22 +000019
20The module standardizes a core set of fast, memory efficient tools that are
Raymond Hettingerf76b9202009-02-17 20:00:59 +000021useful by themselves or in combination. Together, they form an "iterator
22algebra" making it possible to construct specialized tools succinctly and
23efficiently in pure Python.
Georg Brandl116aa622007-08-15 14:28:22 +000024
25For instance, SML provides a tabulation tool: ``tabulate(f)`` which produces a
Ezio Melottib6605992010-01-21 20:57:24 +000026sequence ``f(0), f(1), ...``. The same effect can be achieved in Python
Raymond Hettingera6c60372008-03-13 01:26:19 +000027by combining :func:`map` and :func:`count` to form ``map(f, count())``.
Georg Brandl116aa622007-08-15 14:28:22 +000028
Raymond Hettinger2c109ab2009-03-12 00:29:44 +000029These tools and their built-in counterparts also work well with the high-speed
30functions in the :mod:`operator` module. For example, the multiplication
31operator can be mapped across two vectors to form an efficient dot-product:
32``sum(map(operator.mul, vector1, vector2))``.
Georg Brandl116aa622007-08-15 14:28:22 +000033
Georg Brandl116aa622007-08-15 14:28:22 +000034
Raymond Hettinger6693d7a2017-12-15 13:17:55 -080035**Infinite iterators:**
Georg Brandl116aa622007-08-15 14:28:22 +000036
Raymond Hettinger5bfd8ce2009-04-10 19:02:36 +000037================== ================= ================================================= =========================================
38Iterator Arguments Results Example
39================== ================= ================================================= =========================================
40:func:`count` start, [step] start, start+step, start+2*step, ... ``count(10) --> 10 11 12 13 14 ...``
41:func:`cycle` p p0, p1, ... plast, p0, p1, ... ``cycle('ABCD') --> A B C D A B C D ...``
42:func:`repeat` elem [,n] elem, elem, elem, ... endlessly or up to n times ``repeat(10, 3) --> 10 10 10``
43================== ================= ================================================= =========================================
Georg Brandl116aa622007-08-15 14:28:22 +000044
Raymond Hettingerf76b9202009-02-17 20:00:59 +000045**Iterators terminating on the shortest input sequence:**
46
Benjamin Peterson2989f582014-01-16 10:10:13 -050047============================ ============================ ================================================= =============================================================
48Iterator Arguments Results Example
49============================ ============================ ================================================= =============================================================
50:func:`accumulate` p [,func] p0, p0+p1, p0+p1+p2, ... ``accumulate([1,2,3,4,5]) --> 1 3 6 10 15``
51:func:`chain` p, q, ... p0, p1, ... plast, q0, q1, ... ``chain('ABC', 'DEF') --> A B C D E F``
52:func:`chain.from_iterable` iterable p0, p1, ... plast, q0, q1, ... ``chain.from_iterable(['ABC', 'DEF']) --> A B C D E F``
53:func:`compress` data, selectors (d[0] if s[0]), (d[1] if s[1]), ... ``compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F``
54:func:`dropwhile` pred, seq seq[n], seq[n+1], starting when pred fails ``dropwhile(lambda x: x<5, [1,4,6,4,1]) --> 6 4 1``
55:func:`filterfalse` pred, seq elements of seq where pred(elem) is false ``filterfalse(lambda x: x%2, range(10)) --> 0 2 4 6 8``
Raymond Hettinger49392c62017-09-25 01:21:06 -070056:func:`groupby` iterable[, key] sub-iterators grouped by value of key(v)
Benjamin Peterson2989f582014-01-16 10:10:13 -050057:func:`islice` seq, [start,] stop [, step] elements from seq[start:stop:step] ``islice('ABCDEFG', 2, None) --> C D E F G``
58:func:`starmap` func, seq func(\*seq[0]), func(\*seq[1]), ... ``starmap(pow, [(2,5), (3,2), (10,3)]) --> 32 9 1000``
59:func:`takewhile` pred, seq seq[0], seq[1], until pred fails ``takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4``
60:func:`tee` it, n it1, it2, ... itn splits one iterator into n
61:func:`zip_longest` p, q, ... (p[0], q[0]), (p[1], q[1]), ... ``zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-``
62============================ ============================ ================================================= =============================================================
Raymond Hettingerf76b9202009-02-17 20:00:59 +000063
Raymond Hettinger6693d7a2017-12-15 13:17:55 -080064**Combinatoric iterators:**
Raymond Hettingerf76b9202009-02-17 20:00:59 +000065
Raymond Hettinger7f587cd2009-04-10 19:43:50 +000066============================================== ==================== =============================================================
67Iterator Arguments Results
68============================================== ==================== =============================================================
69:func:`product` p, q, ... [repeat=1] cartesian product, equivalent to a nested for-loop
70:func:`permutations` p[, r] r-length tuples, all possible orderings, no repeated elements
Raymond Hettinger36c3c022009-11-19 01:20:07 +000071:func:`combinations` p, r r-length tuples, in sorted order, no repeated elements
72:func:`combinations_with_replacement` p, r r-length tuples, in sorted order, with repeated elements
Raymond Hettinger7f587cd2009-04-10 19:43:50 +000073``product('ABCD', repeat=2)`` ``AA AB AC AD BA BB BC BD CA CB CC CD DA DB DC DD``
74``permutations('ABCD', 2)`` ``AB AC AD BA BC BD CA CB CD DA DB DC``
75``combinations('ABCD', 2)`` ``AB AC AD BC BD CD``
76``combinations_with_replacement('ABCD', 2)`` ``AA AB AC AD BB BC BD CC CD DD``
77============================================== ==================== =============================================================
Georg Brandl116aa622007-08-15 14:28:22 +000078
79
80.. _itertools-functions:
81
82Itertool functions
83------------------
84
85The following module functions all construct and return iterators. Some provide
86streams of infinite length, so they should only be accessed by functions or
87loops that truncate the stream.
88
Raymond Hettinger5d446132011-03-27 18:52:10 -070089.. function:: accumulate(iterable[, func])
Raymond Hettingeradb81462010-12-01 22:50:36 +000090
Andrew Kuchling15b04eb2014-04-15 22:28:40 -040091 Make an iterator that returns accumulated sums, or accumulated
92 results of other binary functions (specified via the optional
93 *func* argument). If *func* is supplied, it should be a function
94 of two arguments. Elements of the input *iterable* may be any type
95 that can be accepted as arguments to *func*. (For example, with
96 the default operation of addition, elements may be any addable
97 type including :class:`~decimal.Decimal` or
98 :class:`~fractions.Fraction`.) If the input iterable is empty, the
99 output iterable will also be empty.
Raymond Hettingeradb81462010-12-01 22:50:36 +0000100
Raymond Hettinger672866d2016-05-28 00:17:54 -0700101 Roughly equivalent to::
Raymond Hettinger5d446132011-03-27 18:52:10 -0700102
103 def accumulate(iterable, func=operator.add):
Raymond Hettingeradb81462010-12-01 22:50:36 +0000104 'Return running totals'
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000105 # accumulate([1,2,3,4,5]) --> 1 3 6 10 15
Raymond Hettinger5d446132011-03-27 18:52:10 -0700106 # accumulate([1,2,3,4,5], operator.mul) --> 1 2 6 24 120
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000107 it = iter(iterable)
Raymond Hettinger828d9322014-11-22 21:56:23 -0800108 try:
109 total = next(it)
110 except StopIteration:
111 return
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000112 yield total
113 for element in it:
Raymond Hettinger5d446132011-03-27 18:52:10 -0700114 total = func(total, element)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000115 yield total
Raymond Hettingeradb81462010-12-01 22:50:36 +0000116
Raymond Hettinger295c1d42011-04-21 11:09:28 -0700117 There are a number of uses for the *func* argument. It can be set to
118 :func:`min` for a running minimum, :func:`max` for a running maximum, or
119 :func:`operator.mul` for a running product. Amortization tables can be
120 built by accumulating interest and applying payments. First-order
Georg Brandl5d941342016-02-26 19:37:12 +0100121 `recurrence relations <https://en.wikipedia.org/wiki/Recurrence_relation>`_
Raymond Hettinger295c1d42011-04-21 11:09:28 -0700122 can be modeled by supplying the initial value in the iterable and using only
123 the accumulated total in *func* argument::
Raymond Hettinger5d446132011-03-27 18:52:10 -0700124
125 >>> data = [3, 4, 6, 2, 1, 9, 0, 7, 5, 8]
126 >>> list(accumulate(data, operator.mul)) # running product
127 [3, 12, 72, 144, 144, 1296, 0, 0, 0, 0]
128 >>> list(accumulate(data, max)) # running maximum
129 [3, 4, 6, 6, 6, 9, 9, 9, 9, 9]
130
131 # Amortize a 5% loan of 1000 with 4 annual payments of 90
132 >>> cashflows = [1000, -90, -90, -90, -90]
133 >>> list(accumulate(cashflows, lambda bal, pmt: bal*1.05 + pmt))
134 [1000, 960.0, 918.0, 873.9000000000001, 827.5950000000001]
135
Georg Brandl5d941342016-02-26 19:37:12 +0100136 # Chaotic recurrence relation https://en.wikipedia.org/wiki/Logistic_map
Raymond Hettinger295c1d42011-04-21 11:09:28 -0700137 >>> logistic_map = lambda x, _: r * x * (1 - x)
138 >>> r = 3.8
139 >>> x0 = 0.4
140 >>> inputs = repeat(x0, 36) # only the initial value is used
141 >>> [format(x, '.2f') for x in accumulate(inputs, logistic_map)]
142 ['0.40', '0.91', '0.30', '0.81', '0.60', '0.92', '0.29', '0.79', '0.63',
Serhiy Storchakaa4d170d2013-12-23 18:20:51 +0200143 '0.88', '0.39', '0.90', '0.33', '0.84', '0.52', '0.95', '0.18', '0.57',
Raymond Hettinger295c1d42011-04-21 11:09:28 -0700144 '0.93', '0.25', '0.71', '0.79', '0.63', '0.88', '0.39', '0.91', '0.32',
145 '0.83', '0.54', '0.95', '0.20', '0.60', '0.91', '0.30', '0.80', '0.60']
146
Raymond Hettinger64801682013-10-12 16:04:17 -0700147 See :func:`functools.reduce` for a similar function that returns only the
148 final accumulated value.
149
Raymond Hettingeradb81462010-12-01 22:50:36 +0000150 .. versionadded:: 3.2
Georg Brandl116aa622007-08-15 14:28:22 +0000151
Raymond Hettinger5d446132011-03-27 18:52:10 -0700152 .. versionchanged:: 3.3
153 Added the optional *func* parameter.
154
Georg Brandl116aa622007-08-15 14:28:22 +0000155.. function:: chain(*iterables)
156
Raymond Hettinger30c73622010-12-01 23:45:20 +0000157 Make an iterator that returns elements from the first iterable until it is
158 exhausted, then proceeds to the next iterable, until all of the iterables are
159 exhausted. Used for treating consecutive sequences as a single sequence.
Raymond Hettinger672866d2016-05-28 00:17:54 -0700160 Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000161
Raymond Hettinger30c73622010-12-01 23:45:20 +0000162 def chain(*iterables):
163 # chain('ABC', 'DEF') --> A B C D E F
164 for it in iterables:
165 for element in it:
166 yield element
Georg Brandl116aa622007-08-15 14:28:22 +0000167
168
Georg Brandl933b9742010-07-29 14:36:11 +0000169.. classmethod:: chain.from_iterable(iterable)
Christian Heimes70e7ea22008-02-28 20:02:27 +0000170
Raymond Hettinger30c73622010-12-01 23:45:20 +0000171 Alternate constructor for :func:`chain`. Gets chained inputs from a
Raymond Hettinger1e21ebc2013-09-09 01:54:27 -0500172 single iterable argument that is evaluated lazily. Roughly equivalent to::
Christian Heimes70e7ea22008-02-28 20:02:27 +0000173
Raymond Hettinger30c73622010-12-01 23:45:20 +0000174 def from_iterable(iterables):
175 # chain.from_iterable(['ABC', 'DEF']) --> A B C D E F
176 for it in iterables:
177 for element in it:
178 yield element
Christian Heimes70e7ea22008-02-28 20:02:27 +0000179
Christian Heimes78644762008-03-04 23:39:23 +0000180
Christian Heimes836baa52008-02-26 08:18:30 +0000181.. function:: combinations(iterable, r)
182
Raymond Hettinger30c73622010-12-01 23:45:20 +0000183 Return *r* length subsequences of elements from the input *iterable*.
Christian Heimes836baa52008-02-26 08:18:30 +0000184
Raymond Hettinger30c73622010-12-01 23:45:20 +0000185 Combinations are emitted in lexicographic sort order. So, if the
186 input *iterable* is sorted, the combination tuples will be produced
187 in sorted order.
Christian Heimes836baa52008-02-26 08:18:30 +0000188
Raymond Hettinger30c73622010-12-01 23:45:20 +0000189 Elements are treated as unique based on their position, not on their
190 value. So if the input elements are unique, there will be no repeat
191 values in each combination.
Christian Heimes836baa52008-02-26 08:18:30 +0000192
Raymond Hettinger672866d2016-05-28 00:17:54 -0700193 Roughly equivalent to::
Christian Heimes836baa52008-02-26 08:18:30 +0000194
195 def combinations(iterable, r):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000196 # combinations('ABCD', 2) --> AB AC AD BC BD CD
197 # combinations(range(4), 3) --> 012 013 023 123
Christian Heimes836baa52008-02-26 08:18:30 +0000198 pool = tuple(iterable)
Christian Heimes380f7f22008-02-28 11:19:05 +0000199 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000200 if r > n:
201 return
202 indices = list(range(r))
Christian Heimesb558a2e2008-03-02 22:46:37 +0000203 yield tuple(pool[i] for i in indices)
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000204 while True:
Christian Heimes380f7f22008-02-28 11:19:05 +0000205 for i in reversed(range(r)):
Christian Heimesb558a2e2008-03-02 22:46:37 +0000206 if indices[i] != i + n - r:
Christian Heimes836baa52008-02-26 08:18:30 +0000207 break
Christian Heimes380f7f22008-02-28 11:19:05 +0000208 else:
209 return
Christian Heimesb558a2e2008-03-02 22:46:37 +0000210 indices[i] += 1
Christian Heimes380f7f22008-02-28 11:19:05 +0000211 for j in range(i+1, r):
Christian Heimesb558a2e2008-03-02 22:46:37 +0000212 indices[j] = indices[j-1] + 1
213 yield tuple(pool[i] for i in indices)
Christian Heimes836baa52008-02-26 08:18:30 +0000214
Raymond Hettinger30c73622010-12-01 23:45:20 +0000215 The code for :func:`combinations` can be also expressed as a subsequence
216 of :func:`permutations` after filtering entries where the elements are not
217 in sorted order (according to their position in the input pool)::
Christian Heimes78644762008-03-04 23:39:23 +0000218
219 def combinations(iterable, r):
220 pool = tuple(iterable)
221 n = len(pool)
222 for indices in permutations(range(n), r):
223 if sorted(indices) == list(indices):
224 yield tuple(pool[i] for i in indices)
225
Raymond Hettinger30c73622010-12-01 23:45:20 +0000226 The number of items returned is ``n! / r! / (n-r)!`` when ``0 <= r <= n``
227 or zero when ``r > n``.
Christian Heimes836baa52008-02-26 08:18:30 +0000228
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000229.. function:: combinations_with_replacement(iterable, r)
230
Raymond Hettinger30c73622010-12-01 23:45:20 +0000231 Return *r* length subsequences of elements from the input *iterable*
232 allowing individual elements to be repeated more than once.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000233
Raymond Hettinger30c73622010-12-01 23:45:20 +0000234 Combinations are emitted in lexicographic sort order. So, if the
235 input *iterable* is sorted, the combination tuples will be produced
236 in sorted order.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000237
Raymond Hettinger30c73622010-12-01 23:45:20 +0000238 Elements are treated as unique based on their position, not on their
239 value. So if the input elements are unique, the generated combinations
240 will also be unique.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000241
Raymond Hettinger672866d2016-05-28 00:17:54 -0700242 Roughly equivalent to::
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000243
244 def combinations_with_replacement(iterable, r):
245 # combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC
246 pool = tuple(iterable)
247 n = len(pool)
248 if not n and r:
249 return
250 indices = [0] * r
251 yield tuple(pool[i] for i in indices)
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000252 while True:
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000253 for i in reversed(range(r)):
254 if indices[i] != n - 1:
255 break
256 else:
257 return
258 indices[i:] = [indices[i] + 1] * (r - i)
259 yield tuple(pool[i] for i in indices)
260
Raymond Hettinger30c73622010-12-01 23:45:20 +0000261 The code for :func:`combinations_with_replacement` can be also expressed as
262 a subsequence of :func:`product` after filtering entries where the elements
263 are not in sorted order (according to their position in the input pool)::
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000264
265 def combinations_with_replacement(iterable, r):
266 pool = tuple(iterable)
267 n = len(pool)
268 for indices in product(range(n), repeat=r):
269 if sorted(indices) == list(indices):
270 yield tuple(pool[i] for i in indices)
271
Raymond Hettinger30c73622010-12-01 23:45:20 +0000272 The number of items returned is ``(n+r-1)! / r! / (n-1)!`` when ``n > 0``.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000273
Raymond Hettinger30c73622010-12-01 23:45:20 +0000274 .. versionadded:: 3.1
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000275
Georg Brandl67b21b72010-08-17 15:07:14 +0000276
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000277.. function:: compress(data, selectors)
278
Raymond Hettinger30c73622010-12-01 23:45:20 +0000279 Make an iterator that filters elements from *data* returning only those that
280 have a corresponding element in *selectors* that evaluates to ``True``.
281 Stops when either the *data* or *selectors* iterables has been exhausted.
Raymond Hettinger672866d2016-05-28 00:17:54 -0700282 Roughly equivalent to::
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000283
Raymond Hettinger30c73622010-12-01 23:45:20 +0000284 def compress(data, selectors):
285 # compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F
286 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000287
Raymond Hettinger30c73622010-12-01 23:45:20 +0000288 .. versionadded:: 3.1
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000289
290
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000291.. function:: count(start=0, step=1)
Georg Brandl116aa622007-08-15 14:28:22 +0000292
Andrew Kuchlingedb42602013-06-21 08:00:58 -0400293 Make an iterator that returns evenly spaced values starting with number *start*. Often
Raymond Hettinger30c73622010-12-01 23:45:20 +0000294 used as an argument to :func:`map` to generate consecutive data points.
Raymond Hettinger672866d2016-05-28 00:17:54 -0700295 Also, used with :func:`zip` to add sequence numbers. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000296
Raymond Hettinger30c73622010-12-01 23:45:20 +0000297 def count(start=0, step=1):
298 # count(10) --> 10 11 12 13 14 ...
Georg Brandl37a80dc2011-01-13 07:31:18 +0000299 # count(2.5, 0.5) -> 2.5 3.0 3.5 ...
Raymond Hettinger30c73622010-12-01 23:45:20 +0000300 n = start
301 while True:
302 yield n
303 n += step
Georg Brandl116aa622007-08-15 14:28:22 +0000304
Raymond Hettinger30c73622010-12-01 23:45:20 +0000305 When counting with floating point numbers, better accuracy can sometimes be
306 achieved by substituting multiplicative code such as: ``(start + step * i
307 for i in count())``.
Raymond Hettinger5bc472a2009-06-17 01:40:52 +0000308
Raymond Hettinger30c73622010-12-01 23:45:20 +0000309 .. versionchanged:: 3.1
310 Added *step* argument and allowed non-integer arguments.
Georg Brandl116aa622007-08-15 14:28:22 +0000311
312.. function:: cycle(iterable)
313
Raymond Hettinger30c73622010-12-01 23:45:20 +0000314 Make an iterator returning elements from the iterable and saving a copy of each.
315 When the iterable is exhausted, return elements from the saved copy. Repeats
Raymond Hettinger672866d2016-05-28 00:17:54 -0700316 indefinitely. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000317
Raymond Hettinger30c73622010-12-01 23:45:20 +0000318 def cycle(iterable):
319 # cycle('ABCD') --> A B C D A B C D A B C D ...
320 saved = []
321 for element in iterable:
322 yield element
323 saved.append(element)
324 while saved:
325 for element in saved:
Georg Brandl116aa622007-08-15 14:28:22 +0000326 yield element
327
Raymond Hettinger30c73622010-12-01 23:45:20 +0000328 Note, this member of the toolkit may require significant auxiliary storage
329 (depending on the length of the iterable).
Georg Brandl116aa622007-08-15 14:28:22 +0000330
331
332.. function:: dropwhile(predicate, iterable)
333
Raymond Hettinger30c73622010-12-01 23:45:20 +0000334 Make an iterator that drops elements from the iterable as long as the predicate
335 is true; afterwards, returns every element. Note, the iterator does not produce
336 *any* output until the predicate first becomes false, so it may have a lengthy
Raymond Hettinger672866d2016-05-28 00:17:54 -0700337 start-up time. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000338
Raymond Hettinger30c73622010-12-01 23:45:20 +0000339 def dropwhile(predicate, iterable):
340 # dropwhile(lambda x: x<5, [1,4,6,4,1]) --> 6 4 1
341 iterable = iter(iterable)
342 for x in iterable:
343 if not predicate(x):
344 yield x
345 break
346 for x in iterable:
347 yield x
Georg Brandl116aa622007-08-15 14:28:22 +0000348
Raymond Hettinger749761e2009-01-27 04:42:48 +0000349.. function:: filterfalse(predicate, iterable)
350
Raymond Hettinger30c73622010-12-01 23:45:20 +0000351 Make an iterator that filters elements from iterable returning only those for
352 which the predicate is ``False``. If *predicate* is ``None``, return the items
Raymond Hettinger672866d2016-05-28 00:17:54 -0700353 that are false. Roughly equivalent to::
Raymond Hettinger749761e2009-01-27 04:42:48 +0000354
Raymond Hettinger30c73622010-12-01 23:45:20 +0000355 def filterfalse(predicate, iterable):
356 # filterfalse(lambda x: x%2, range(10)) --> 0 2 4 6 8
357 if predicate is None:
358 predicate = bool
359 for x in iterable:
360 if not predicate(x):
361 yield x
Raymond Hettinger749761e2009-01-27 04:42:48 +0000362
Georg Brandl116aa622007-08-15 14:28:22 +0000363
Georg Brandl3dd33882009-06-01 17:35:27 +0000364.. function:: groupby(iterable, key=None)
Georg Brandl116aa622007-08-15 14:28:22 +0000365
Raymond Hettinger30c73622010-12-01 23:45:20 +0000366 Make an iterator that returns consecutive keys and groups from the *iterable*.
367 The *key* is a function computing a key value for each element. If not
368 specified or is ``None``, *key* defaults to an identity function and returns
369 the element unchanged. Generally, the iterable needs to already be sorted on
370 the same key function.
Georg Brandl116aa622007-08-15 14:28:22 +0000371
Raymond Hettinger30c73622010-12-01 23:45:20 +0000372 The operation of :func:`groupby` is similar to the ``uniq`` filter in Unix. It
373 generates a break or new group every time the value of the key function changes
374 (which is why it is usually necessary to have sorted the data using the same key
375 function). That behavior differs from SQL's GROUP BY which aggregates common
376 elements regardless of their input order.
Georg Brandl116aa622007-08-15 14:28:22 +0000377
Raymond Hettinger30c73622010-12-01 23:45:20 +0000378 The returned group is itself an iterator that shares the underlying iterable
379 with :func:`groupby`. Because the source is shared, when the :func:`groupby`
380 object is advanced, the previous group is no longer visible. So, if that data
381 is needed later, it should be stored as a list::
Georg Brandl116aa622007-08-15 14:28:22 +0000382
Raymond Hettinger30c73622010-12-01 23:45:20 +0000383 groups = []
384 uniquekeys = []
385 data = sorted(data, key=keyfunc)
386 for k, g in groupby(data, keyfunc):
387 groups.append(list(g)) # Store group iterator as a list
388 uniquekeys.append(k)
Georg Brandl116aa622007-08-15 14:28:22 +0000389
Raymond Hettinger672866d2016-05-28 00:17:54 -0700390 :func:`groupby` is roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000391
Raymond Hettinger30c73622010-12-01 23:45:20 +0000392 class groupby:
393 # [k for k, g in groupby('AAAABBBCCDAABBB')] --> A B C D A B
394 # [list(g) for k, g in groupby('AAAABBBCCD')] --> AAAA BBB CC D
395 def __init__(self, iterable, key=None):
396 if key is None:
397 key = lambda x: x
398 self.keyfunc = key
399 self.it = iter(iterable)
400 self.tgtkey = self.currkey = self.currvalue = object()
401 def __iter__(self):
402 return self
403 def __next__(self):
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300404 self.id = object()
Raymond Hettinger30c73622010-12-01 23:45:20 +0000405 while self.currkey == self.tgtkey:
406 self.currvalue = next(self.it) # Exit on StopIteration
407 self.currkey = self.keyfunc(self.currvalue)
408 self.tgtkey = self.currkey
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300409 return (self.currkey, self._grouper(self.tgtkey, self.id))
410 def _grouper(self, tgtkey, id):
411 while self.id is id and self.currkey == tgtkey:
Raymond Hettinger30c73622010-12-01 23:45:20 +0000412 yield self.currvalue
Raymond Hettinger828d9322014-11-22 21:56:23 -0800413 try:
414 self.currvalue = next(self.it)
415 except StopIteration:
416 return
Raymond Hettinger30c73622010-12-01 23:45:20 +0000417 self.currkey = self.keyfunc(self.currvalue)
Georg Brandl116aa622007-08-15 14:28:22 +0000418
Georg Brandl116aa622007-08-15 14:28:22 +0000419
Ezio Melottie0add762012-09-14 06:32:35 +0300420.. function:: islice(iterable, stop)
421 islice(iterable, start, stop[, step])
Georg Brandl116aa622007-08-15 14:28:22 +0000422
Raymond Hettinger30c73622010-12-01 23:45:20 +0000423 Make an iterator that returns selected elements from the iterable. If *start* is
424 non-zero, then elements from the iterable are skipped until start is reached.
425 Afterward, elements are returned consecutively unless *step* is set higher than
426 one which results in items being skipped. If *stop* is ``None``, then iteration
427 continues until the iterator is exhausted, if at all; otherwise, it stops at the
428 specified position. Unlike regular slicing, :func:`islice` does not support
429 negative values for *start*, *stop*, or *step*. Can be used to extract related
430 fields from data where the internal structure has been flattened (for example, a
Raymond Hettinger672866d2016-05-28 00:17:54 -0700431 multi-line report may list a name field on every third line). Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000432
Raymond Hettinger30c73622010-12-01 23:45:20 +0000433 def islice(iterable, *args):
434 # islice('ABCDEFG', 2) --> A B
435 # islice('ABCDEFG', 2, 4) --> C D
436 # islice('ABCDEFG', 2, None) --> C D E F G
437 # islice('ABCDEFG', 0, None, 2) --> A C E G
438 s = slice(*args)
Cheryl Sabellada1734c2018-03-26 21:29:33 -0400439 start, stop, step = s.start or 0, s.stop or sys.maxsize, s.step or 1
440 it = iter(range(start, stop, step))
Raymond Hettinger828d9322014-11-22 21:56:23 -0800441 try:
442 nexti = next(it)
443 except StopIteration:
Cheryl Sabellada1734c2018-03-26 21:29:33 -0400444 # Consume *iterable* up to the *start* position.
445 for i, element in zip(range(start), iterable):
446 pass
Raymond Hettinger828d9322014-11-22 21:56:23 -0800447 return
Cheryl Sabellada1734c2018-03-26 21:29:33 -0400448 try:
449 for i, element in enumerate(iterable):
450 if i == nexti:
451 yield element
452 nexti = next(it)
453 except StopIteration:
454 # Consume to *stop*.
455 for i, element in zip(range(i + 1, stop), iterable):
456 pass
Georg Brandl116aa622007-08-15 14:28:22 +0000457
Raymond Hettinger30c73622010-12-01 23:45:20 +0000458 If *start* is ``None``, then iteration starts at zero. If *step* is ``None``,
459 then the step defaults to one.
Georg Brandl116aa622007-08-15 14:28:22 +0000460
Georg Brandl116aa622007-08-15 14:28:22 +0000461
Georg Brandl3dd33882009-06-01 17:35:27 +0000462.. function:: permutations(iterable, r=None)
Christian Heimes70e7ea22008-02-28 20:02:27 +0000463
Raymond Hettinger30c73622010-12-01 23:45:20 +0000464 Return successive *r* length permutations of elements in the *iterable*.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000465
Raymond Hettinger30c73622010-12-01 23:45:20 +0000466 If *r* is not specified or is ``None``, then *r* defaults to the length
467 of the *iterable* and all possible full-length permutations
468 are generated.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000469
Raymond Hettinger30c73622010-12-01 23:45:20 +0000470 Permutations are emitted in lexicographic sort order. So, if the
471 input *iterable* is sorted, the permutation tuples will be produced
472 in sorted order.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000473
Raymond Hettinger30c73622010-12-01 23:45:20 +0000474 Elements are treated as unique based on their position, not on their
475 value. So if the input elements are unique, there will be no repeat
476 values in each permutation.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000477
Raymond Hettinger672866d2016-05-28 00:17:54 -0700478 Roughly equivalent to::
Christian Heimesb558a2e2008-03-02 22:46:37 +0000479
480 def permutations(iterable, r=None):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000481 # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
482 # permutations(range(3)) --> 012 021 102 120 201 210
Christian Heimesb558a2e2008-03-02 22:46:37 +0000483 pool = tuple(iterable)
484 n = len(pool)
485 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000486 if r > n:
487 return
488 indices = list(range(n))
Sandro Tosi73866622011-12-25 17:25:45 +0100489 cycles = list(range(n, n-r, -1))
Christian Heimesb558a2e2008-03-02 22:46:37 +0000490 yield tuple(pool[i] for i in indices[:r])
491 while n:
492 for i in reversed(range(r)):
493 cycles[i] -= 1
494 if cycles[i] == 0:
495 indices[i:] = indices[i+1:] + indices[i:i+1]
496 cycles[i] = n - i
497 else:
498 j = cycles[i]
499 indices[i], indices[-j] = indices[-j], indices[i]
500 yield tuple(pool[i] for i in indices[:r])
501 break
502 else:
503 return
Christian Heimes70e7ea22008-02-28 20:02:27 +0000504
Raymond Hettinger30c73622010-12-01 23:45:20 +0000505 The code for :func:`permutations` can be also expressed as a subsequence of
506 :func:`product`, filtered to exclude entries with repeated elements (those
507 from the same position in the input pool)::
Christian Heimes78644762008-03-04 23:39:23 +0000508
509 def permutations(iterable, r=None):
510 pool = tuple(iterable)
511 n = len(pool)
512 r = n if r is None else r
513 for indices in product(range(n), repeat=r):
514 if len(set(indices)) == r:
515 yield tuple(pool[i] for i in indices)
516
Raymond Hettinger30c73622010-12-01 23:45:20 +0000517 The number of items returned is ``n! / (n-r)!`` when ``0 <= r <= n``
518 or zero when ``r > n``.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000519
Georg Brandl3dd33882009-06-01 17:35:27 +0000520.. function:: product(*iterables, repeat=1)
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000521
Raymond Hettinger30c73622010-12-01 23:45:20 +0000522 Cartesian product of input iterables.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000523
Raymond Hettinger672866d2016-05-28 00:17:54 -0700524 Roughly equivalent to nested for-loops in a generator expression. For example,
Raymond Hettinger30c73622010-12-01 23:45:20 +0000525 ``product(A, B)`` returns the same as ``((x,y) for x in A for y in B)``.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000526
Raymond Hettinger30c73622010-12-01 23:45:20 +0000527 The nested loops cycle like an odometer with the rightmost element advancing
528 on every iteration. This pattern creates a lexicographic ordering so that if
529 the input's iterables are sorted, the product tuples are emitted in sorted
530 order.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000531
Raymond Hettinger30c73622010-12-01 23:45:20 +0000532 To compute the product of an iterable with itself, specify the number of
533 repetitions with the optional *repeat* keyword argument. For example,
534 ``product(A, repeat=4)`` means the same as ``product(A, A, A, A)``.
Christian Heimes9e7f1d22008-02-28 12:27:11 +0000535
Raymond Hettinger672866d2016-05-28 00:17:54 -0700536 This function is roughly equivalent to the following code, except that the
Raymond Hettinger30c73622010-12-01 23:45:20 +0000537 actual implementation does not build up intermediate results in memory::
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000538
Raymond Hettinger30c73622010-12-01 23:45:20 +0000539 def product(*args, repeat=1):
540 # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy
541 # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111
542 pools = [tuple(pool) for pool in args] * repeat
543 result = [[]]
544 for pool in pools:
545 result = [x+[y] for x in result for y in pool]
546 for prod in result:
547 yield tuple(prod)
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000548
549
Raymond Hettingerd75ad442009-06-01 19:16:52 +0000550.. function:: repeat(object[, times])
Georg Brandl116aa622007-08-15 14:28:22 +0000551
Raymond Hettinger30c73622010-12-01 23:45:20 +0000552 Make an iterator that returns *object* over and over again. Runs indefinitely
553 unless the *times* argument is specified. Used as argument to :func:`map` for
554 invariant parameters to the called function. Also used with :func:`zip` to
Raymond Hettinger672866d2016-05-28 00:17:54 -0700555 create an invariant part of a tuple record.
556
557 Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000558
Raymond Hettinger30c73622010-12-01 23:45:20 +0000559 def repeat(object, times=None):
560 # repeat(10, 3) --> 10 10 10
561 if times is None:
562 while True:
563 yield object
564 else:
565 for i in range(times):
566 yield object
Georg Brandl116aa622007-08-15 14:28:22 +0000567
Raymond Hettingerfc3ba6b2012-02-01 09:07:40 -0800568 A common use for *repeat* is to supply a stream of constant values to *map*
569 or *zip*::
570
571 >>> list(map(pow, range(10), repeat(2)))
572 [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
Georg Brandl116aa622007-08-15 14:28:22 +0000573
574.. function:: starmap(function, iterable)
575
Raymond Hettinger30c73622010-12-01 23:45:20 +0000576 Make an iterator that computes the function using arguments obtained from
577 the iterable. Used instead of :func:`map` when argument parameters are already
578 grouped in tuples from a single iterable (the data has been "pre-zipped"). The
579 difference between :func:`map` and :func:`starmap` parallels the distinction
Raymond Hettinger672866d2016-05-28 00:17:54 -0700580 between ``function(a,b)`` and ``function(*c)``. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000581
Raymond Hettinger30c73622010-12-01 23:45:20 +0000582 def starmap(function, iterable):
583 # starmap(pow, [(2,5), (3,2), (10,3)]) --> 32 9 1000
584 for args in iterable:
585 yield function(*args)
Christian Heimes679db4a2008-01-18 09:56:22 +0000586
Georg Brandl116aa622007-08-15 14:28:22 +0000587
588.. function:: takewhile(predicate, iterable)
589
Raymond Hettinger30c73622010-12-01 23:45:20 +0000590 Make an iterator that returns elements from the iterable as long as the
Raymond Hettinger672866d2016-05-28 00:17:54 -0700591 predicate is true. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000592
Raymond Hettinger30c73622010-12-01 23:45:20 +0000593 def takewhile(predicate, iterable):
594 # takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4
595 for x in iterable:
596 if predicate(x):
597 yield x
598 else:
599 break
Georg Brandl116aa622007-08-15 14:28:22 +0000600
601
Georg Brandl3dd33882009-06-01 17:35:27 +0000602.. function:: tee(iterable, n=2)
Georg Brandl116aa622007-08-15 14:28:22 +0000603
Raymond Hettingerab425aa2016-04-26 00:10:00 -0700604 Return *n* independent iterators from a single iterable.
605
606 The following Python code helps explain what *tee* does (although the actual
Serhiy Storchaka4ecfa452016-05-16 09:31:54 +0300607 implementation is more complex and uses only a single underlying
Raymond Hettinger672866d2016-05-28 00:17:54 -0700608 :abbr:`FIFO (first-in, first-out)` queue).
609
610 Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000611
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000612 def tee(iterable, n=2):
613 it = iter(iterable)
614 deques = [collections.deque() for i in range(n)]
615 def gen(mydeque):
616 while True:
617 if not mydeque: # when the local deque is empty
Raymond Hettinger828d9322014-11-22 21:56:23 -0800618 try:
619 newval = next(it) # fetch a new value and
620 except StopIteration:
621 return
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000622 for d in deques: # load it to all the deques
623 d.append(newval)
624 yield mydeque.popleft()
625 return tuple(gen(d) for d in deques)
Georg Brandl116aa622007-08-15 14:28:22 +0000626
Raymond Hettinger30c73622010-12-01 23:45:20 +0000627 Once :func:`tee` has made a split, the original *iterable* should not be
628 used anywhere else; otherwise, the *iterable* could get advanced without
629 the tee objects being informed.
Georg Brandl116aa622007-08-15 14:28:22 +0000630
Raymond Hettinger30c73622010-12-01 23:45:20 +0000631 This itertool may require significant auxiliary storage (depending on how
632 much temporary data needs to be stored). In general, if one iterator uses
633 most or all of the data before another iterator starts, it is faster to use
634 :func:`list` instead of :func:`tee`.
Georg Brandl116aa622007-08-15 14:28:22 +0000635
Georg Brandl116aa622007-08-15 14:28:22 +0000636
Georg Brandl3dd33882009-06-01 17:35:27 +0000637.. function:: zip_longest(*iterables, fillvalue=None)
Raymond Hettinger749761e2009-01-27 04:42:48 +0000638
Raymond Hettinger30c73622010-12-01 23:45:20 +0000639 Make an iterator that aggregates elements from each of the iterables. If the
640 iterables are of uneven length, missing values are filled-in with *fillvalue*.
Raymond Hettinger672866d2016-05-28 00:17:54 -0700641 Iteration continues until the longest iterable is exhausted. Roughly equivalent to::
Raymond Hettinger749761e2009-01-27 04:42:48 +0000642
Raymond Hettinger3147b042017-09-07 14:01:44 -0700643 def zip_longest(*args, fillvalue=None):
Raymond Hettinger30c73622010-12-01 23:45:20 +0000644 # zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-
Raymond Hettinger3147b042017-09-07 14:01:44 -0700645 iterators = [iter(it) for it in args]
646 num_active = len(iterators)
647 if not num_active:
648 return
649 while True:
650 values = []
651 for i, it in enumerate(iterators):
652 try:
653 value = next(it)
654 except StopIteration:
655 num_active -= 1
656 if not num_active:
657 return
658 iterators[i] = repeat(fillvalue)
659 value = fillvalue
660 values.append(value)
661 yield tuple(values)
Raymond Hettinger749761e2009-01-27 04:42:48 +0000662
Raymond Hettinger30c73622010-12-01 23:45:20 +0000663 If one of the iterables is potentially infinite, then the :func:`zip_longest`
664 function should be wrapped with something that limits the number of calls
665 (for example :func:`islice` or :func:`takewhile`). If not specified,
666 *fillvalue* defaults to ``None``.
Raymond Hettinger749761e2009-01-27 04:42:48 +0000667
668
Georg Brandl116aa622007-08-15 14:28:22 +0000669.. _itertools-recipes:
670
Raymond Hettinger1fa76822010-12-06 23:31:36 +0000671Itertools Recipes
672-----------------
Georg Brandl116aa622007-08-15 14:28:22 +0000673
674This section shows recipes for creating an extended toolset using the existing
675itertools as building blocks.
676
677The extended tools offer the same high performance as the underlying toolset.
678The superior memory performance is kept by processing elements one at a time
679rather than bringing the whole iterable into memory all at once. Code volume is
680kept small by linking the tools together in a functional style which helps
681eliminate temporary variables. High speed is retained by preferring
Georg Brandl9afde1c2007-11-01 20:32:30 +0000682"vectorized" building blocks over the use of for-loops and :term:`generator`\s
Christian Heimesfe337bf2008-03-23 21:54:12 +0000683which incur interpreter overhead.
684
685.. testcode::
Georg Brandl116aa622007-08-15 14:28:22 +0000686
Raymond Hettinger30c73622010-12-01 23:45:20 +0000687 def take(n, iterable):
688 "Return first n items of the iterable as a list"
689 return list(islice(iterable, n))
Georg Brandl116aa622007-08-15 14:28:22 +0000690
Raymond Hettinger9265dd72018-04-08 08:44:20 -0700691 def prepend(value, iterator):
692 "Prepend a single value in front of an iterator"
693 # prepend(1, [2, 3, 4]) -> 1 2 3 4
694 return chain([value], iterator)
695
Raymond Hettinger30c73622010-12-01 23:45:20 +0000696 def tabulate(function, start=0):
697 "Return function(0), function(1), ..."
698 return map(function, count(start))
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000699
Raymond Hettingerdfe098d2014-05-25 22:03:56 -0700700 def tail(n, iterable):
701 "Return an iterator over the last n items"
702 # tail(3, 'ABCDEFG') --> E F G
703 return iter(collections.deque(iterable, maxlen=n))
704
Cheryl Sabellada1734c2018-03-26 21:29:33 -0400705 def consume(iterator, n=None):
706 "Advance the iterator n-steps ahead. If n is None, consume entirely."
Raymond Hettinger30c73622010-12-01 23:45:20 +0000707 # Use functions that consume iterators at C speed.
708 if n is None:
709 # feed the entire iterator into a zero-length deque
710 collections.deque(iterator, maxlen=0)
711 else:
712 # advance to the empty slice starting at position n
713 next(islice(iterator, n, n), None)
Raymond Hettingerfa007962009-03-09 11:55:25 +0000714
Raymond Hettinger30c73622010-12-01 23:45:20 +0000715 def nth(iterable, n, default=None):
716 "Returns the nth item or a default value"
717 return next(islice(iterable, n, None), default)
Georg Brandl116aa622007-08-15 14:28:22 +0000718
Raymond Hettingere525ee32016-03-06 18:11:38 -0800719 def all_equal(iterable):
720 "Returns True if all the elements are equal to each other"
721 g = groupby(iterable)
722 return next(g, True) and not next(g, False)
723
Raymond Hettinger30c73622010-12-01 23:45:20 +0000724 def quantify(iterable, pred=bool):
725 "Count how many times the predicate is true"
726 return sum(map(pred, iterable))
Georg Brandl116aa622007-08-15 14:28:22 +0000727
Raymond Hettinger30c73622010-12-01 23:45:20 +0000728 def padnone(iterable):
729 """Returns the sequence elements and then returns None indefinitely.
Georg Brandl116aa622007-08-15 14:28:22 +0000730
Raymond Hettinger30c73622010-12-01 23:45:20 +0000731 Useful for emulating the behavior of the built-in map() function.
732 """
733 return chain(iterable, repeat(None))
Georg Brandl116aa622007-08-15 14:28:22 +0000734
Raymond Hettinger30c73622010-12-01 23:45:20 +0000735 def ncycles(iterable, n):
736 "Returns the sequence elements n times"
737 return chain.from_iterable(repeat(tuple(iterable), n))
Georg Brandl116aa622007-08-15 14:28:22 +0000738
Raymond Hettinger30c73622010-12-01 23:45:20 +0000739 def dotproduct(vec1, vec2):
740 return sum(map(operator.mul, vec1, vec2))
Georg Brandl116aa622007-08-15 14:28:22 +0000741
Raymond Hettinger30c73622010-12-01 23:45:20 +0000742 def flatten(listOfLists):
743 "Flatten one level of nesting"
744 return chain.from_iterable(listOfLists)
Georg Brandl116aa622007-08-15 14:28:22 +0000745
Raymond Hettinger30c73622010-12-01 23:45:20 +0000746 def repeatfunc(func, times=None, *args):
747 """Repeat calls to func with specified arguments.
Georg Brandl116aa622007-08-15 14:28:22 +0000748
Raymond Hettinger30c73622010-12-01 23:45:20 +0000749 Example: repeatfunc(random.random)
750 """
751 if times is None:
752 return starmap(func, repeat(args))
753 return starmap(func, repeat(args, times))
Georg Brandl116aa622007-08-15 14:28:22 +0000754
Raymond Hettinger30c73622010-12-01 23:45:20 +0000755 def pairwise(iterable):
756 "s -> (s0,s1), (s1,s2), (s2, s3), ..."
757 a, b = tee(iterable)
758 next(b, None)
759 return zip(a, b)
Georg Brandl116aa622007-08-15 14:28:22 +0000760
Raymond Hettinger44571da2013-05-05 19:53:41 -0700761 def grouper(iterable, n, fillvalue=None):
Raymond Hettinger9ae94732012-07-08 16:04:03 -0700762 "Collect data into fixed-length chunks or blocks"
Raymond Hettinger44571da2013-05-05 19:53:41 -0700763 # grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx"
Raymond Hettinger30c73622010-12-01 23:45:20 +0000764 args = [iter(iterable)] * n
765 return zip_longest(*args, fillvalue=fillvalue)
Georg Brandl116aa622007-08-15 14:28:22 +0000766
Raymond Hettinger30c73622010-12-01 23:45:20 +0000767 def roundrobin(*iterables):
768 "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
769 # Recipe credited to George Sakkis
Raymond Hettinger337cbba2017-11-21 00:23:34 -0800770 num_active = len(iterables)
Raymond Hettinger30c73622010-12-01 23:45:20 +0000771 nexts = cycle(iter(it).__next__ for it in iterables)
Raymond Hettinger337cbba2017-11-21 00:23:34 -0800772 while num_active:
Raymond Hettinger30c73622010-12-01 23:45:20 +0000773 try:
774 for next in nexts:
775 yield next()
776 except StopIteration:
Raymond Hettinger337cbba2017-11-21 00:23:34 -0800777 # Remove the iterator we just exhausted from the cycle.
778 num_active -= 1
779 nexts = cycle(islice(nexts, num_active))
Georg Brandl116aa622007-08-15 14:28:22 +0000780
Raymond Hettinger30c73622010-12-01 23:45:20 +0000781 def partition(pred, iterable):
782 'Use a predicate to partition entries into false entries and true entries'
783 # partition(is_odd, range(10)) --> 0 2 4 6 8 and 1 3 5 7 9
784 t1, t2 = tee(iterable)
785 return filterfalse(pred, t1), filter(pred, t2)
Raymond Hettinger5ce0aa22010-12-01 10:49:19 +0000786
Raymond Hettinger30c73622010-12-01 23:45:20 +0000787 def powerset(iterable):
788 "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
789 s = list(iterable)
790 return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettinger08d01ee2010-08-07 05:36:53 +0000791
Raymond Hettinger30c73622010-12-01 23:45:20 +0000792 def unique_everseen(iterable, key=None):
793 "List unique elements, preserving order. Remember all elements ever seen."
794 # unique_everseen('AAAABBBCCDAABBB') --> A B C D
795 # unique_everseen('ABBCcAD', str.lower) --> A B C D
796 seen = set()
797 seen_add = seen.add
798 if key is None:
799 for element in filterfalse(seen.__contains__, iterable):
800 seen_add(element)
801 yield element
802 else:
803 for element in iterable:
804 k = key(element)
805 if k not in seen:
806 seen_add(k)
807 yield element
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000808
Raymond Hettinger30c73622010-12-01 23:45:20 +0000809 def unique_justseen(iterable, key=None):
810 "List unique elements, preserving order. Remember only the element just seen."
811 # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
812 # unique_justseen('ABBCcAD', str.lower) --> A B C A D
813 return map(next, map(itemgetter(1), groupby(iterable, key)))
Raymond Hettingerad9d96b2009-01-02 21:39:07 +0000814
Raymond Hettinger30c73622010-12-01 23:45:20 +0000815 def iter_except(func, exception, first=None):
816 """ Call a function repeatedly until an exception is raised.
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000817
Raymond Hettinger30c73622010-12-01 23:45:20 +0000818 Converts a call-until-exception interface to an iterator interface.
Andrew Kuchling1d7d5802013-06-20 21:17:41 -0400819 Like builtins.iter(func, sentinel) but uses an exception instead
Raymond Hettinger30c73622010-12-01 23:45:20 +0000820 of a sentinel to end the loop.
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000821
Raymond Hettinger30c73622010-12-01 23:45:20 +0000822 Examples:
823 iter_except(functools.partial(heappop, h), IndexError) # priority queue iterator
824 iter_except(d.popitem, KeyError) # non-blocking dict iterator
825 iter_except(d.popleft, IndexError) # non-blocking deque iterator
826 iter_except(q.get_nowait, Queue.Empty) # loop over a producer Queue
827 iter_except(s.pop, KeyError) # non-blocking set iterator
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000828
Raymond Hettinger30c73622010-12-01 23:45:20 +0000829 """
830 try:
831 if first is not None:
832 yield first() # For database APIs needing an initial cast to db.first()
Raymond Hettingera503f702016-03-13 00:12:31 -0800833 while True:
Raymond Hettinger30c73622010-12-01 23:45:20 +0000834 yield func()
835 except exception:
836 pass
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000837
Raymond Hettinger31b26f62014-04-02 03:16:42 -0700838 def first_true(iterable, default=False, pred=None):
839 """Returns the first true value in the iterable.
840
841 If no true value is found, returns *default*
842
843 If *pred* is not None, returns the first item
844 for which pred(item) is true.
845
846 """
847 # first_true([a,b,c], x) --> a or b or c or x
848 # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
849 return next(filter(pred, iterable), default)
850
Raymond Hettinger30c73622010-12-01 23:45:20 +0000851 def random_product(*args, repeat=1):
852 "Random selection from itertools.product(*args, **kwds)"
853 pools = [tuple(pool) for pool in args] * repeat
854 return tuple(random.choice(pool) for pool in pools)
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000855
Raymond Hettinger30c73622010-12-01 23:45:20 +0000856 def random_permutation(iterable, r=None):
857 "Random selection from itertools.permutations(iterable, r)"
858 pool = tuple(iterable)
859 r = len(pool) if r is None else r
860 return tuple(random.sample(pool, r))
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000861
Raymond Hettinger30c73622010-12-01 23:45:20 +0000862 def random_combination(iterable, r):
863 "Random selection from itertools.combinations(iterable, r)"
864 pool = tuple(iterable)
865 n = len(pool)
866 indices = sorted(random.sample(range(n), r))
867 return tuple(pool[i] for i in indices)
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000868
Raymond Hettinger30c73622010-12-01 23:45:20 +0000869 def random_combination_with_replacement(iterable, r):
870 "Random selection from itertools.combinations_with_replacement(iterable, r)"
871 pool = tuple(iterable)
872 n = len(pool)
873 indices = sorted(random.randrange(n) for i in range(r))
874 return tuple(pool[i] for i in indices)
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000875
Raymond Hettingerd37258d2018-01-13 10:35:40 -0800876 def nth_combination(iterable, r, index):
877 'Equivalent to list(combinations(iterable, r))[index]'
878 pool = tuple(iterable)
879 n = len(pool)
880 if r < 0 or r > n:
881 raise ValueError
882 c = 1
883 k = min(r, n-r)
884 for i in range(1, k+1):
885 c = c * (n - k + i) // i
886 if index < 0:
887 index += c
888 if index < 0 or index >= c:
889 raise IndexError
890 result = []
891 while r:
892 c, n, r = c*r//n, n-1, r-1
893 while index >= c:
894 index -= c
895 c, n = c*(n-r)//n, n-1
896 result.append(pool[-1-n])
897 return tuple(result)
898
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000899Note, many of the above recipes can be optimized by replacing global lookups
900with local variables defined as default values. For example, the
901*dotproduct* recipe can be written as::
902
Raymond Hettinger30c73622010-12-01 23:45:20 +0000903 def dotproduct(vec1, vec2, sum=sum, map=map, mul=operator.mul):
904 return sum(map(mul, vec1, vec2))