blob: c6e2544d3c401a0ee5f716a584a47ed0962cd1ba [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.
Georg Brandl116aa622007-08-15 14:28:22 +00006.. moduleauthor:: Raymond Hettinger <python@rcn.com>
7.. sectionauthor:: Raymond Hettinger <python@rcn.com>
8
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
14
Raymond Hettingerf76b9202009-02-17 20:00:59 +000015This module implements a number of :term:`iterator` building blocks inspired
16by constructs from APL, Haskell, and SML. Each has been recast in a form
17suitable for Python.
Georg Brandl116aa622007-08-15 14:28:22 +000018
19The module standardizes a core set of fast, memory efficient tools that are
Raymond Hettingerf76b9202009-02-17 20:00:59 +000020useful by themselves or in combination. Together, they form an "iterator
21algebra" making it possible to construct specialized tools succinctly and
22efficiently in pure Python.
Georg Brandl116aa622007-08-15 14:28:22 +000023
24For instance, SML provides a tabulation tool: ``tabulate(f)`` which produces a
Ezio Melottib6605992010-01-21 20:57:24 +000025sequence ``f(0), f(1), ...``. The same effect can be achieved in Python
Raymond Hettingera6c60372008-03-13 01:26:19 +000026by combining :func:`map` and :func:`count` to form ``map(f, count())``.
Georg Brandl116aa622007-08-15 14:28:22 +000027
Raymond Hettinger2c109ab2009-03-12 00:29:44 +000028These tools and their built-in counterparts also work well with the high-speed
29functions in the :mod:`operator` module. For example, the multiplication
30operator can be mapped across two vectors to form an efficient dot-product:
31``sum(map(operator.mul, vector1, vector2))``.
Georg Brandl116aa622007-08-15 14:28:22 +000032
Georg Brandl116aa622007-08-15 14:28:22 +000033
Raymond Hettingerf76b9202009-02-17 20:00:59 +000034**Infinite Iterators:**
Georg Brandl116aa622007-08-15 14:28:22 +000035
Raymond Hettinger5bfd8ce2009-04-10 19:02:36 +000036================== ================= ================================================= =========================================
37Iterator Arguments Results Example
38================== ================= ================================================= =========================================
39:func:`count` start, [step] start, start+step, start+2*step, ... ``count(10) --> 10 11 12 13 14 ...``
40:func:`cycle` p p0, p1, ... plast, p0, p1, ... ``cycle('ABCD') --> A B C D A B C D ...``
41:func:`repeat` elem [,n] elem, elem, elem, ... endlessly or up to n times ``repeat(10, 3) --> 10 10 10``
42================== ================= ================================================= =========================================
Georg Brandl116aa622007-08-15 14:28:22 +000043
Raymond Hettingerf76b9202009-02-17 20:00:59 +000044**Iterators terminating on the shortest input sequence:**
45
Benjamin Peterson2989f582014-01-16 10:10:13 -050046============================ ============================ ================================================= =============================================================
47Iterator Arguments Results Example
48============================ ============================ ================================================= =============================================================
49:func:`accumulate` p [,func] p0, p0+p1, p0+p1+p2, ... ``accumulate([1,2,3,4,5]) --> 1 3 6 10 15``
50:func:`chain` p, q, ... p0, p1, ... plast, q0, q1, ... ``chain('ABC', 'DEF') --> A B C D E F``
51:func:`chain.from_iterable` iterable p0, p1, ... plast, q0, q1, ... ``chain.from_iterable(['ABC', 'DEF']) --> A B C D E F``
52: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``
53: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``
54:func:`filterfalse` pred, seq elements of seq where pred(elem) is false ``filterfalse(lambda x: x%2, range(10)) --> 0 2 4 6 8``
55:func:`groupby` iterable[, keyfunc] sub-iterators grouped by value of keyfunc(v)
56:func:`islice` seq, [start,] stop [, step] elements from seq[start:stop:step] ``islice('ABCDEFG', 2, None) --> C D E F G``
57:func:`starmap` func, seq func(\*seq[0]), func(\*seq[1]), ... ``starmap(pow, [(2,5), (3,2), (10,3)]) --> 32 9 1000``
58:func:`takewhile` pred, seq seq[0], seq[1], until pred fails ``takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4``
59:func:`tee` it, n it1, it2, ... itn splits one iterator into n
60:func:`zip_longest` p, q, ... (p[0], q[0]), (p[1], q[1]), ... ``zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-``
61============================ ============================ ================================================= =============================================================
Raymond Hettingerf76b9202009-02-17 20:00:59 +000062
63**Combinatoric generators:**
64
Raymond Hettinger7f587cd2009-04-10 19:43:50 +000065============================================== ==================== =============================================================
66Iterator Arguments Results
67============================================== ==================== =============================================================
68:func:`product` p, q, ... [repeat=1] cartesian product, equivalent to a nested for-loop
69:func:`permutations` p[, r] r-length tuples, all possible orderings, no repeated elements
Raymond Hettinger36c3c022009-11-19 01:20:07 +000070:func:`combinations` p, r r-length tuples, in sorted order, no repeated elements
71:func:`combinations_with_replacement` p, r r-length tuples, in sorted order, with repeated elements
Raymond Hettinger7f587cd2009-04-10 19:43:50 +000072``product('ABCD', repeat=2)`` ``AA AB AC AD BA BB BC BD CA CB CC CD DA DB DC DD``
73``permutations('ABCD', 2)`` ``AB AC AD BA BC BD CA CB CD DA DB DC``
74``combinations('ABCD', 2)`` ``AB AC AD BC BD CD``
75``combinations_with_replacement('ABCD', 2)`` ``AA AB AC AD BB BC BD CC CD DD``
76============================================== ==================== =============================================================
Georg Brandl116aa622007-08-15 14:28:22 +000077
78
79.. _itertools-functions:
80
81Itertool functions
82------------------
83
84The following module functions all construct and return iterators. Some provide
85streams of infinite length, so they should only be accessed by functions or
86loops that truncate the stream.
87
Raymond Hettinger5d446132011-03-27 18:52:10 -070088.. function:: accumulate(iterable[, func])
Raymond Hettingeradb81462010-12-01 22:50:36 +000089
Andrew Kuchling15b04eb2014-04-15 22:28:40 -040090 Make an iterator that returns accumulated sums, or accumulated
91 results of other binary functions (specified via the optional
92 *func* argument). If *func* is supplied, it should be a function
93 of two arguments. Elements of the input *iterable* may be any type
94 that can be accepted as arguments to *func*. (For example, with
95 the default operation of addition, elements may be any addable
96 type including :class:`~decimal.Decimal` or
97 :class:`~fractions.Fraction`.) If the input iterable is empty, the
98 output iterable will also be empty.
Raymond Hettingeradb81462010-12-01 22:50:36 +000099
Raymond Hettinger819581b2016-05-28 00:10:56 -0700100 Roughly equivalent to::
Raymond Hettinger5d446132011-03-27 18:52:10 -0700101
102 def accumulate(iterable, func=operator.add):
Raymond Hettingeradb81462010-12-01 22:50:36 +0000103 'Return running totals'
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000104 # accumulate([1,2,3,4,5]) --> 1 3 6 10 15
Raymond Hettinger5d446132011-03-27 18:52:10 -0700105 # accumulate([1,2,3,4,5], operator.mul) --> 1 2 6 24 120
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000106 it = iter(iterable)
Raymond Hettinger828d9322014-11-22 21:56:23 -0800107 try:
108 total = next(it)
109 except StopIteration:
110 return
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000111 yield total
112 for element in it:
Raymond Hettinger5d446132011-03-27 18:52:10 -0700113 total = func(total, element)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000114 yield total
Raymond Hettingeradb81462010-12-01 22:50:36 +0000115
Raymond Hettinger295c1d42011-04-21 11:09:28 -0700116 There are a number of uses for the *func* argument. It can be set to
117 :func:`min` for a running minimum, :func:`max` for a running maximum, or
118 :func:`operator.mul` for a running product. Amortization tables can be
119 built by accumulating interest and applying payments. First-order
Georg Brandl5d941342016-02-26 19:37:12 +0100120 `recurrence relations <https://en.wikipedia.org/wiki/Recurrence_relation>`_
Raymond Hettinger295c1d42011-04-21 11:09:28 -0700121 can be modeled by supplying the initial value in the iterable and using only
122 the accumulated total in *func* argument::
Raymond Hettinger5d446132011-03-27 18:52:10 -0700123
124 >>> data = [3, 4, 6, 2, 1, 9, 0, 7, 5, 8]
125 >>> list(accumulate(data, operator.mul)) # running product
126 [3, 12, 72, 144, 144, 1296, 0, 0, 0, 0]
127 >>> list(accumulate(data, max)) # running maximum
128 [3, 4, 6, 6, 6, 9, 9, 9, 9, 9]
129
130 # Amortize a 5% loan of 1000 with 4 annual payments of 90
131 >>> cashflows = [1000, -90, -90, -90, -90]
132 >>> list(accumulate(cashflows, lambda bal, pmt: bal*1.05 + pmt))
133 [1000, 960.0, 918.0, 873.9000000000001, 827.5950000000001]
134
Georg Brandl5d941342016-02-26 19:37:12 +0100135 # Chaotic recurrence relation https://en.wikipedia.org/wiki/Logistic_map
Raymond Hettinger295c1d42011-04-21 11:09:28 -0700136 >>> logistic_map = lambda x, _: r * x * (1 - x)
137 >>> r = 3.8
138 >>> x0 = 0.4
139 >>> inputs = repeat(x0, 36) # only the initial value is used
140 >>> [format(x, '.2f') for x in accumulate(inputs, logistic_map)]
141 ['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 +0200142 '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 -0700143 '0.93', '0.25', '0.71', '0.79', '0.63', '0.88', '0.39', '0.91', '0.32',
144 '0.83', '0.54', '0.95', '0.20', '0.60', '0.91', '0.30', '0.80', '0.60']
145
Raymond Hettinger64801682013-10-12 16:04:17 -0700146 See :func:`functools.reduce` for a similar function that returns only the
147 final accumulated value.
148
Raymond Hettingeradb81462010-12-01 22:50:36 +0000149 .. versionadded:: 3.2
Georg Brandl116aa622007-08-15 14:28:22 +0000150
Raymond Hettinger5d446132011-03-27 18:52:10 -0700151 .. versionchanged:: 3.3
152 Added the optional *func* parameter.
153
Georg Brandl116aa622007-08-15 14:28:22 +0000154.. function:: chain(*iterables)
155
Raymond Hettinger30c73622010-12-01 23:45:20 +0000156 Make an iterator that returns elements from the first iterable until it is
157 exhausted, then proceeds to the next iterable, until all of the iterables are
158 exhausted. Used for treating consecutive sequences as a single sequence.
Raymond Hettinger819581b2016-05-28 00:10:56 -0700159 Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000160
Raymond Hettinger30c73622010-12-01 23:45:20 +0000161 def chain(*iterables):
162 # chain('ABC', 'DEF') --> A B C D E F
163 for it in iterables:
164 for element in it:
165 yield element
Georg Brandl116aa622007-08-15 14:28:22 +0000166
167
Georg Brandl933b9742010-07-29 14:36:11 +0000168.. classmethod:: chain.from_iterable(iterable)
Christian Heimes70e7ea22008-02-28 20:02:27 +0000169
Raymond Hettinger30c73622010-12-01 23:45:20 +0000170 Alternate constructor for :func:`chain`. Gets chained inputs from a
Raymond Hettinger1e21ebc2013-09-09 01:54:27 -0500171 single iterable argument that is evaluated lazily. Roughly equivalent to::
Christian Heimes70e7ea22008-02-28 20:02:27 +0000172
Raymond Hettinger30c73622010-12-01 23:45:20 +0000173 def from_iterable(iterables):
174 # chain.from_iterable(['ABC', 'DEF']) --> A B C D E F
175 for it in iterables:
176 for element in it:
177 yield element
Christian Heimes70e7ea22008-02-28 20:02:27 +0000178
Christian Heimes78644762008-03-04 23:39:23 +0000179
Christian Heimes836baa52008-02-26 08:18:30 +0000180.. function:: combinations(iterable, r)
181
Raymond Hettinger30c73622010-12-01 23:45:20 +0000182 Return *r* length subsequences of elements from the input *iterable*.
Christian Heimes836baa52008-02-26 08:18:30 +0000183
Raymond Hettinger30c73622010-12-01 23:45:20 +0000184 Combinations are emitted in lexicographic sort order. So, if the
185 input *iterable* is sorted, the combination tuples will be produced
186 in sorted order.
Christian Heimes836baa52008-02-26 08:18:30 +0000187
Raymond Hettinger30c73622010-12-01 23:45:20 +0000188 Elements are treated as unique based on their position, not on their
189 value. So if the input elements are unique, there will be no repeat
190 values in each combination.
Christian Heimes836baa52008-02-26 08:18:30 +0000191
Raymond Hettinger819581b2016-05-28 00:10:56 -0700192 Roughly equivalent to::
Christian Heimes836baa52008-02-26 08:18:30 +0000193
194 def combinations(iterable, r):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000195 # combinations('ABCD', 2) --> AB AC AD BC BD CD
196 # combinations(range(4), 3) --> 012 013 023 123
Christian Heimes836baa52008-02-26 08:18:30 +0000197 pool = tuple(iterable)
Christian Heimes380f7f22008-02-28 11:19:05 +0000198 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000199 if r > n:
200 return
201 indices = list(range(r))
Christian Heimesb558a2e2008-03-02 22:46:37 +0000202 yield tuple(pool[i] for i in indices)
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000203 while True:
Christian Heimes380f7f22008-02-28 11:19:05 +0000204 for i in reversed(range(r)):
Christian Heimesb558a2e2008-03-02 22:46:37 +0000205 if indices[i] != i + n - r:
Christian Heimes836baa52008-02-26 08:18:30 +0000206 break
Christian Heimes380f7f22008-02-28 11:19:05 +0000207 else:
208 return
Christian Heimesb558a2e2008-03-02 22:46:37 +0000209 indices[i] += 1
Christian Heimes380f7f22008-02-28 11:19:05 +0000210 for j in range(i+1, r):
Christian Heimesb558a2e2008-03-02 22:46:37 +0000211 indices[j] = indices[j-1] + 1
212 yield tuple(pool[i] for i in indices)
Christian Heimes836baa52008-02-26 08:18:30 +0000213
Raymond Hettinger30c73622010-12-01 23:45:20 +0000214 The code for :func:`combinations` can be also expressed as a subsequence
215 of :func:`permutations` after filtering entries where the elements are not
216 in sorted order (according to their position in the input pool)::
Christian Heimes78644762008-03-04 23:39:23 +0000217
218 def combinations(iterable, r):
219 pool = tuple(iterable)
220 n = len(pool)
221 for indices in permutations(range(n), r):
222 if sorted(indices) == list(indices):
223 yield tuple(pool[i] for i in indices)
224
Raymond Hettinger30c73622010-12-01 23:45:20 +0000225 The number of items returned is ``n! / r! / (n-r)!`` when ``0 <= r <= n``
226 or zero when ``r > n``.
Christian Heimes836baa52008-02-26 08:18:30 +0000227
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000228.. function:: combinations_with_replacement(iterable, r)
229
Raymond Hettinger30c73622010-12-01 23:45:20 +0000230 Return *r* length subsequences of elements from the input *iterable*
231 allowing individual elements to be repeated more than once.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000232
Raymond Hettinger30c73622010-12-01 23:45:20 +0000233 Combinations are emitted in lexicographic sort order. So, if the
234 input *iterable* is sorted, the combination tuples will be produced
235 in sorted order.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000236
Raymond Hettinger30c73622010-12-01 23:45:20 +0000237 Elements are treated as unique based on their position, not on their
238 value. So if the input elements are unique, the generated combinations
239 will also be unique.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000240
Raymond Hettinger819581b2016-05-28 00:10:56 -0700241 Roughly equivalent to::
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000242
243 def combinations_with_replacement(iterable, r):
244 # combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC
245 pool = tuple(iterable)
246 n = len(pool)
247 if not n and r:
248 return
249 indices = [0] * r
250 yield tuple(pool[i] for i in indices)
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000251 while True:
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000252 for i in reversed(range(r)):
253 if indices[i] != n - 1:
254 break
255 else:
256 return
257 indices[i:] = [indices[i] + 1] * (r - i)
258 yield tuple(pool[i] for i in indices)
259
Raymond Hettinger30c73622010-12-01 23:45:20 +0000260 The code for :func:`combinations_with_replacement` can be also expressed as
261 a subsequence of :func:`product` after filtering entries where the elements
262 are not in sorted order (according to their position in the input pool)::
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000263
264 def combinations_with_replacement(iterable, r):
265 pool = tuple(iterable)
266 n = len(pool)
267 for indices in product(range(n), repeat=r):
268 if sorted(indices) == list(indices):
269 yield tuple(pool[i] for i in indices)
270
Raymond Hettinger30c73622010-12-01 23:45:20 +0000271 The number of items returned is ``(n+r-1)! / r! / (n-1)!`` when ``n > 0``.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000272
Raymond Hettinger30c73622010-12-01 23:45:20 +0000273 .. versionadded:: 3.1
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000274
Georg Brandl67b21b72010-08-17 15:07:14 +0000275
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000276.. function:: compress(data, selectors)
277
Raymond Hettinger30c73622010-12-01 23:45:20 +0000278 Make an iterator that filters elements from *data* returning only those that
279 have a corresponding element in *selectors* that evaluates to ``True``.
280 Stops when either the *data* or *selectors* iterables has been exhausted.
Raymond Hettinger819581b2016-05-28 00:10:56 -0700281 Roughly equivalent to::
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000282
Raymond Hettinger30c73622010-12-01 23:45:20 +0000283 def compress(data, selectors):
284 # compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F
285 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000286
Raymond Hettinger30c73622010-12-01 23:45:20 +0000287 .. versionadded:: 3.1
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000288
289
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000290.. function:: count(start=0, step=1)
Georg Brandl116aa622007-08-15 14:28:22 +0000291
Andrew Kuchlingedb42602013-06-21 08:00:58 -0400292 Make an iterator that returns evenly spaced values starting with number *start*. Often
Raymond Hettinger30c73622010-12-01 23:45:20 +0000293 used as an argument to :func:`map` to generate consecutive data points.
Raymond Hettinger819581b2016-05-28 00:10:56 -0700294 Also, used with :func:`zip` to add sequence numbers. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000295
Raymond Hettinger30c73622010-12-01 23:45:20 +0000296 def count(start=0, step=1):
297 # count(10) --> 10 11 12 13 14 ...
Georg Brandl37a80dc2011-01-13 07:31:18 +0000298 # count(2.5, 0.5) -> 2.5 3.0 3.5 ...
Raymond Hettinger30c73622010-12-01 23:45:20 +0000299 n = start
300 while True:
301 yield n
302 n += step
Georg Brandl116aa622007-08-15 14:28:22 +0000303
Raymond Hettinger30c73622010-12-01 23:45:20 +0000304 When counting with floating point numbers, better accuracy can sometimes be
305 achieved by substituting multiplicative code such as: ``(start + step * i
306 for i in count())``.
Raymond Hettinger5bc472a2009-06-17 01:40:52 +0000307
Raymond Hettinger30c73622010-12-01 23:45:20 +0000308 .. versionchanged:: 3.1
309 Added *step* argument and allowed non-integer arguments.
Georg Brandl116aa622007-08-15 14:28:22 +0000310
311.. function:: cycle(iterable)
312
Raymond Hettinger30c73622010-12-01 23:45:20 +0000313 Make an iterator returning elements from the iterable and saving a copy of each.
314 When the iterable is exhausted, return elements from the saved copy. Repeats
Raymond Hettinger819581b2016-05-28 00:10:56 -0700315 indefinitely. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000316
Raymond Hettinger30c73622010-12-01 23:45:20 +0000317 def cycle(iterable):
318 # cycle('ABCD') --> A B C D A B C D A B C D ...
319 saved = []
320 for element in iterable:
321 yield element
322 saved.append(element)
323 while saved:
324 for element in saved:
Georg Brandl116aa622007-08-15 14:28:22 +0000325 yield element
326
Raymond Hettinger30c73622010-12-01 23:45:20 +0000327 Note, this member of the toolkit may require significant auxiliary storage
328 (depending on the length of the iterable).
Georg Brandl116aa622007-08-15 14:28:22 +0000329
330
331.. function:: dropwhile(predicate, iterable)
332
Raymond Hettinger30c73622010-12-01 23:45:20 +0000333 Make an iterator that drops elements from the iterable as long as the predicate
334 is true; afterwards, returns every element. Note, the iterator does not produce
335 *any* output until the predicate first becomes false, so it may have a lengthy
Raymond Hettinger819581b2016-05-28 00:10:56 -0700336 start-up time. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000337
Raymond Hettinger30c73622010-12-01 23:45:20 +0000338 def dropwhile(predicate, iterable):
339 # dropwhile(lambda x: x<5, [1,4,6,4,1]) --> 6 4 1
340 iterable = iter(iterable)
341 for x in iterable:
342 if not predicate(x):
343 yield x
344 break
345 for x in iterable:
346 yield x
Georg Brandl116aa622007-08-15 14:28:22 +0000347
Raymond Hettinger749761e2009-01-27 04:42:48 +0000348.. function:: filterfalse(predicate, iterable)
349
Raymond Hettinger30c73622010-12-01 23:45:20 +0000350 Make an iterator that filters elements from iterable returning only those for
351 which the predicate is ``False``. If *predicate* is ``None``, return the items
Raymond Hettinger819581b2016-05-28 00:10:56 -0700352 that are false. Roughly equivalent to::
Raymond Hettinger749761e2009-01-27 04:42:48 +0000353
Raymond Hettinger30c73622010-12-01 23:45:20 +0000354 def filterfalse(predicate, iterable):
355 # filterfalse(lambda x: x%2, range(10)) --> 0 2 4 6 8
356 if predicate is None:
357 predicate = bool
358 for x in iterable:
359 if not predicate(x):
360 yield x
Raymond Hettinger749761e2009-01-27 04:42:48 +0000361
Georg Brandl116aa622007-08-15 14:28:22 +0000362
Georg Brandl3dd33882009-06-01 17:35:27 +0000363.. function:: groupby(iterable, key=None)
Georg Brandl116aa622007-08-15 14:28:22 +0000364
Raymond Hettinger30c73622010-12-01 23:45:20 +0000365 Make an iterator that returns consecutive keys and groups from the *iterable*.
366 The *key* is a function computing a key value for each element. If not
367 specified or is ``None``, *key* defaults to an identity function and returns
368 the element unchanged. Generally, the iterable needs to already be sorted on
369 the same key function.
Georg Brandl116aa622007-08-15 14:28:22 +0000370
Raymond Hettinger30c73622010-12-01 23:45:20 +0000371 The operation of :func:`groupby` is similar to the ``uniq`` filter in Unix. It
372 generates a break or new group every time the value of the key function changes
373 (which is why it is usually necessary to have sorted the data using the same key
374 function). That behavior differs from SQL's GROUP BY which aggregates common
375 elements regardless of their input order.
Georg Brandl116aa622007-08-15 14:28:22 +0000376
Raymond Hettinger30c73622010-12-01 23:45:20 +0000377 The returned group is itself an iterator that shares the underlying iterable
378 with :func:`groupby`. Because the source is shared, when the :func:`groupby`
379 object is advanced, the previous group is no longer visible. So, if that data
380 is needed later, it should be stored as a list::
Georg Brandl116aa622007-08-15 14:28:22 +0000381
Raymond Hettinger30c73622010-12-01 23:45:20 +0000382 groups = []
383 uniquekeys = []
384 data = sorted(data, key=keyfunc)
385 for k, g in groupby(data, keyfunc):
386 groups.append(list(g)) # Store group iterator as a list
387 uniquekeys.append(k)
Georg Brandl116aa622007-08-15 14:28:22 +0000388
Raymond Hettinger819581b2016-05-28 00:10:56 -0700389 :func:`groupby` is roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000390
Raymond Hettinger30c73622010-12-01 23:45:20 +0000391 class groupby:
392 # [k for k, g in groupby('AAAABBBCCDAABBB')] --> A B C D A B
393 # [list(g) for k, g in groupby('AAAABBBCCD')] --> AAAA BBB CC D
394 def __init__(self, iterable, key=None):
395 if key is None:
396 key = lambda x: x
397 self.keyfunc = key
398 self.it = iter(iterable)
399 self.tgtkey = self.currkey = self.currvalue = object()
400 def __iter__(self):
401 return self
402 def __next__(self):
403 while self.currkey == self.tgtkey:
404 self.currvalue = next(self.it) # Exit on StopIteration
405 self.currkey = self.keyfunc(self.currvalue)
406 self.tgtkey = self.currkey
407 return (self.currkey, self._grouper(self.tgtkey))
408 def _grouper(self, tgtkey):
409 while self.currkey == tgtkey:
410 yield self.currvalue
Raymond Hettinger828d9322014-11-22 21:56:23 -0800411 try:
412 self.currvalue = next(self.it)
413 except StopIteration:
414 return
Raymond Hettinger30c73622010-12-01 23:45:20 +0000415 self.currkey = self.keyfunc(self.currvalue)
Georg Brandl116aa622007-08-15 14:28:22 +0000416
Georg Brandl116aa622007-08-15 14:28:22 +0000417
Ezio Melottie0add762012-09-14 06:32:35 +0300418.. function:: islice(iterable, stop)
419 islice(iterable, start, stop[, step])
Georg Brandl116aa622007-08-15 14:28:22 +0000420
Raymond Hettinger30c73622010-12-01 23:45:20 +0000421 Make an iterator that returns selected elements from the iterable. If *start* is
422 non-zero, then elements from the iterable are skipped until start is reached.
423 Afterward, elements are returned consecutively unless *step* is set higher than
424 one which results in items being skipped. If *stop* is ``None``, then iteration
425 continues until the iterator is exhausted, if at all; otherwise, it stops at the
426 specified position. Unlike regular slicing, :func:`islice` does not support
427 negative values for *start*, *stop*, or *step*. Can be used to extract related
428 fields from data where the internal structure has been flattened (for example, a
Raymond Hettinger819581b2016-05-28 00:10:56 -0700429 multi-line report may list a name field on every third line). Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000430
Raymond Hettinger30c73622010-12-01 23:45:20 +0000431 def islice(iterable, *args):
432 # islice('ABCDEFG', 2) --> A B
433 # islice('ABCDEFG', 2, 4) --> C D
434 # islice('ABCDEFG', 2, None) --> C D E F G
435 # islice('ABCDEFG', 0, None, 2) --> A C E G
436 s = slice(*args)
437 it = iter(range(s.start or 0, s.stop or sys.maxsize, s.step or 1))
Raymond Hettinger828d9322014-11-22 21:56:23 -0800438 try:
439 nexti = next(it)
440 except StopIteration:
441 return
Raymond Hettinger30c73622010-12-01 23:45:20 +0000442 for i, element in enumerate(iterable):
443 if i == nexti:
444 yield element
445 nexti = next(it)
Georg Brandl116aa622007-08-15 14:28:22 +0000446
Raymond Hettinger30c73622010-12-01 23:45:20 +0000447 If *start* is ``None``, then iteration starts at zero. If *step* is ``None``,
448 then the step defaults to one.
Georg Brandl116aa622007-08-15 14:28:22 +0000449
Georg Brandl116aa622007-08-15 14:28:22 +0000450
Georg Brandl3dd33882009-06-01 17:35:27 +0000451.. function:: permutations(iterable, r=None)
Christian Heimes70e7ea22008-02-28 20:02:27 +0000452
Raymond Hettinger30c73622010-12-01 23:45:20 +0000453 Return successive *r* length permutations of elements in the *iterable*.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000454
Raymond Hettinger30c73622010-12-01 23:45:20 +0000455 If *r* is not specified or is ``None``, then *r* defaults to the length
456 of the *iterable* and all possible full-length permutations
457 are generated.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000458
Raymond Hettinger30c73622010-12-01 23:45:20 +0000459 Permutations are emitted in lexicographic sort order. So, if the
460 input *iterable* is sorted, the permutation tuples will be produced
461 in sorted order.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000462
Raymond Hettinger30c73622010-12-01 23:45:20 +0000463 Elements are treated as unique based on their position, not on their
464 value. So if the input elements are unique, there will be no repeat
465 values in each permutation.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000466
Raymond Hettinger819581b2016-05-28 00:10:56 -0700467 Roughly equivalent to::
Christian Heimesb558a2e2008-03-02 22:46:37 +0000468
469 def permutations(iterable, r=None):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000470 # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
471 # permutations(range(3)) --> 012 021 102 120 201 210
Christian Heimesb558a2e2008-03-02 22:46:37 +0000472 pool = tuple(iterable)
473 n = len(pool)
474 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000475 if r > n:
476 return
477 indices = list(range(n))
Sandro Tosi73866622011-12-25 17:25:45 +0100478 cycles = list(range(n, n-r, -1))
Christian Heimesb558a2e2008-03-02 22:46:37 +0000479 yield tuple(pool[i] for i in indices[:r])
480 while n:
481 for i in reversed(range(r)):
482 cycles[i] -= 1
483 if cycles[i] == 0:
484 indices[i:] = indices[i+1:] + indices[i:i+1]
485 cycles[i] = n - i
486 else:
487 j = cycles[i]
488 indices[i], indices[-j] = indices[-j], indices[i]
489 yield tuple(pool[i] for i in indices[:r])
490 break
491 else:
492 return
Christian Heimes70e7ea22008-02-28 20:02:27 +0000493
Raymond Hettinger30c73622010-12-01 23:45:20 +0000494 The code for :func:`permutations` can be also expressed as a subsequence of
495 :func:`product`, filtered to exclude entries with repeated elements (those
496 from the same position in the input pool)::
Christian Heimes78644762008-03-04 23:39:23 +0000497
498 def permutations(iterable, r=None):
499 pool = tuple(iterable)
500 n = len(pool)
501 r = n if r is None else r
502 for indices in product(range(n), repeat=r):
503 if len(set(indices)) == r:
504 yield tuple(pool[i] for i in indices)
505
Raymond Hettinger30c73622010-12-01 23:45:20 +0000506 The number of items returned is ``n! / (n-r)!`` when ``0 <= r <= n``
507 or zero when ``r > n``.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000508
Georg Brandl3dd33882009-06-01 17:35:27 +0000509.. function:: product(*iterables, repeat=1)
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000510
Raymond Hettinger30c73622010-12-01 23:45:20 +0000511 Cartesian product of input iterables.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000512
Raymond Hettinger819581b2016-05-28 00:10:56 -0700513 Roughly equivalent to nested for-loops in a generator expression. For example,
Raymond Hettinger30c73622010-12-01 23:45:20 +0000514 ``product(A, B)`` returns the same as ``((x,y) for x in A for y in B)``.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000515
Raymond Hettinger30c73622010-12-01 23:45:20 +0000516 The nested loops cycle like an odometer with the rightmost element advancing
517 on every iteration. This pattern creates a lexicographic ordering so that if
518 the input's iterables are sorted, the product tuples are emitted in sorted
519 order.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000520
Raymond Hettinger30c73622010-12-01 23:45:20 +0000521 To compute the product of an iterable with itself, specify the number of
522 repetitions with the optional *repeat* keyword argument. For example,
523 ``product(A, repeat=4)`` means the same as ``product(A, A, A, A)``.
Christian Heimes9e7f1d22008-02-28 12:27:11 +0000524
Raymond Hettinger819581b2016-05-28 00:10:56 -0700525 This function is roughly equivalent to the following code, except that the
Raymond Hettinger30c73622010-12-01 23:45:20 +0000526 actual implementation does not build up intermediate results in memory::
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000527
Raymond Hettinger30c73622010-12-01 23:45:20 +0000528 def product(*args, repeat=1):
529 # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy
530 # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111
531 pools = [tuple(pool) for pool in args] * repeat
532 result = [[]]
533 for pool in pools:
534 result = [x+[y] for x in result for y in pool]
535 for prod in result:
536 yield tuple(prod)
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000537
538
Raymond Hettingerd75ad442009-06-01 19:16:52 +0000539.. function:: repeat(object[, times])
Georg Brandl116aa622007-08-15 14:28:22 +0000540
Raymond Hettinger30c73622010-12-01 23:45:20 +0000541 Make an iterator that returns *object* over and over again. Runs indefinitely
542 unless the *times* argument is specified. Used as argument to :func:`map` for
543 invariant parameters to the called function. Also used with :func:`zip` to
Raymond Hettinger819581b2016-05-28 00:10:56 -0700544 create an invariant part of a tuple record.
545
546 Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000547
Raymond Hettinger30c73622010-12-01 23:45:20 +0000548 def repeat(object, times=None):
549 # repeat(10, 3) --> 10 10 10
550 if times is None:
551 while True:
552 yield object
553 else:
554 for i in range(times):
555 yield object
Georg Brandl116aa622007-08-15 14:28:22 +0000556
Raymond Hettingerfc3ba6b2012-02-01 09:07:40 -0800557 A common use for *repeat* is to supply a stream of constant values to *map*
558 or *zip*::
559
560 >>> list(map(pow, range(10), repeat(2)))
561 [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
Georg Brandl116aa622007-08-15 14:28:22 +0000562
563.. function:: starmap(function, iterable)
564
Raymond Hettinger30c73622010-12-01 23:45:20 +0000565 Make an iterator that computes the function using arguments obtained from
566 the iterable. Used instead of :func:`map` when argument parameters are already
567 grouped in tuples from a single iterable (the data has been "pre-zipped"). The
568 difference between :func:`map` and :func:`starmap` parallels the distinction
Raymond Hettinger819581b2016-05-28 00:10:56 -0700569 between ``function(a,b)`` and ``function(*c)``. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000570
Raymond Hettinger30c73622010-12-01 23:45:20 +0000571 def starmap(function, iterable):
572 # starmap(pow, [(2,5), (3,2), (10,3)]) --> 32 9 1000
573 for args in iterable:
574 yield function(*args)
Christian Heimes679db4a2008-01-18 09:56:22 +0000575
Georg Brandl116aa622007-08-15 14:28:22 +0000576
577.. function:: takewhile(predicate, iterable)
578
Raymond Hettinger30c73622010-12-01 23:45:20 +0000579 Make an iterator that returns elements from the iterable as long as the
Raymond Hettinger819581b2016-05-28 00:10:56 -0700580 predicate is true. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000581
Raymond Hettinger30c73622010-12-01 23:45:20 +0000582 def takewhile(predicate, iterable):
583 # takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4
584 for x in iterable:
585 if predicate(x):
586 yield x
587 else:
588 break
Georg Brandl116aa622007-08-15 14:28:22 +0000589
590
Georg Brandl3dd33882009-06-01 17:35:27 +0000591.. function:: tee(iterable, n=2)
Georg Brandl116aa622007-08-15 14:28:22 +0000592
Raymond Hettinger819581b2016-05-28 00:10:56 -0700593 Return *n* independent iterators from a single iterable. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000594
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000595 def tee(iterable, n=2):
596 it = iter(iterable)
597 deques = [collections.deque() for i in range(n)]
598 def gen(mydeque):
599 while True:
600 if not mydeque: # when the local deque is empty
Raymond Hettinger828d9322014-11-22 21:56:23 -0800601 try:
602 newval = next(it) # fetch a new value and
603 except StopIteration:
604 return
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000605 for d in deques: # load it to all the deques
606 d.append(newval)
607 yield mydeque.popleft()
608 return tuple(gen(d) for d in deques)
Georg Brandl116aa622007-08-15 14:28:22 +0000609
Raymond Hettinger30c73622010-12-01 23:45:20 +0000610 Once :func:`tee` has made a split, the original *iterable* should not be
611 used anywhere else; otherwise, the *iterable* could get advanced without
612 the tee objects being informed.
Georg Brandl116aa622007-08-15 14:28:22 +0000613
Raymond Hettinger30c73622010-12-01 23:45:20 +0000614 This itertool may require significant auxiliary storage (depending on how
615 much temporary data needs to be stored). In general, if one iterator uses
616 most or all of the data before another iterator starts, it is faster to use
617 :func:`list` instead of :func:`tee`.
Georg Brandl116aa622007-08-15 14:28:22 +0000618
Georg Brandl116aa622007-08-15 14:28:22 +0000619
Georg Brandl3dd33882009-06-01 17:35:27 +0000620.. function:: zip_longest(*iterables, fillvalue=None)
Raymond Hettinger749761e2009-01-27 04:42:48 +0000621
Raymond Hettinger30c73622010-12-01 23:45:20 +0000622 Make an iterator that aggregates elements from each of the iterables. If the
623 iterables are of uneven length, missing values are filled-in with *fillvalue*.
Raymond Hettinger819581b2016-05-28 00:10:56 -0700624 Iteration continues until the longest iterable is exhausted. Roughly equivalent to::
Raymond Hettinger749761e2009-01-27 04:42:48 +0000625
Raymond Hettinger6f45d182011-10-30 15:06:14 -0700626 class ZipExhausted(Exception):
627 pass
628
629 def zip_longest(*args, **kwds):
Raymond Hettinger30c73622010-12-01 23:45:20 +0000630 # zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-
Raymond Hettinger6f45d182011-10-30 15:06:14 -0700631 fillvalue = kwds.get('fillvalue')
632 counter = len(args) - 1
633 def sentinel():
634 nonlocal counter
635 if not counter:
636 raise ZipExhausted
637 counter -= 1
638 yield fillvalue
Raymond Hettinger30c73622010-12-01 23:45:20 +0000639 fillers = repeat(fillvalue)
Raymond Hettinger6f45d182011-10-30 15:06:14 -0700640 iterators = [chain(it, sentinel(), fillers) for it in args]
Raymond Hettinger30c73622010-12-01 23:45:20 +0000641 try:
Raymond Hettinger6f45d182011-10-30 15:06:14 -0700642 while iterators:
643 yield tuple(map(next, iterators))
644 except ZipExhausted:
Raymond Hettinger30c73622010-12-01 23:45:20 +0000645 pass
Raymond Hettinger749761e2009-01-27 04:42:48 +0000646
Raymond Hettinger30c73622010-12-01 23:45:20 +0000647 If one of the iterables is potentially infinite, then the :func:`zip_longest`
648 function should be wrapped with something that limits the number of calls
649 (for example :func:`islice` or :func:`takewhile`). If not specified,
650 *fillvalue* defaults to ``None``.
Raymond Hettinger749761e2009-01-27 04:42:48 +0000651
652
Georg Brandl116aa622007-08-15 14:28:22 +0000653.. _itertools-recipes:
654
Raymond Hettinger1fa76822010-12-06 23:31:36 +0000655Itertools Recipes
656-----------------
Georg Brandl116aa622007-08-15 14:28:22 +0000657
658This section shows recipes for creating an extended toolset using the existing
659itertools as building blocks.
660
661The extended tools offer the same high performance as the underlying toolset.
662The superior memory performance is kept by processing elements one at a time
663rather than bringing the whole iterable into memory all at once. Code volume is
664kept small by linking the tools together in a functional style which helps
665eliminate temporary variables. High speed is retained by preferring
Georg Brandl9afde1c2007-11-01 20:32:30 +0000666"vectorized" building blocks over the use of for-loops and :term:`generator`\s
Christian Heimesfe337bf2008-03-23 21:54:12 +0000667which incur interpreter overhead.
668
669.. testcode::
Georg Brandl116aa622007-08-15 14:28:22 +0000670
Raymond Hettinger30c73622010-12-01 23:45:20 +0000671 def take(n, iterable):
672 "Return first n items of the iterable as a list"
673 return list(islice(iterable, n))
Georg Brandl116aa622007-08-15 14:28:22 +0000674
Raymond Hettinger30c73622010-12-01 23:45:20 +0000675 def tabulate(function, start=0):
676 "Return function(0), function(1), ..."
677 return map(function, count(start))
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000678
Raymond Hettingerdfe098d2014-05-25 22:03:56 -0700679 def tail(n, iterable):
680 "Return an iterator over the last n items"
681 # tail(3, 'ABCDEFG') --> E F G
682 return iter(collections.deque(iterable, maxlen=n))
683
Raymond Hettinger30c73622010-12-01 23:45:20 +0000684 def consume(iterator, n):
685 "Advance the iterator n-steps ahead. If n is none, consume entirely."
686 # Use functions that consume iterators at C speed.
687 if n is None:
688 # feed the entire iterator into a zero-length deque
689 collections.deque(iterator, maxlen=0)
690 else:
691 # advance to the empty slice starting at position n
692 next(islice(iterator, n, n), None)
Raymond Hettingerfa007962009-03-09 11:55:25 +0000693
Raymond Hettinger30c73622010-12-01 23:45:20 +0000694 def nth(iterable, n, default=None):
695 "Returns the nth item or a default value"
696 return next(islice(iterable, n, None), default)
Georg Brandl116aa622007-08-15 14:28:22 +0000697
Raymond Hettingere525ee32016-03-06 18:11:38 -0800698 def all_equal(iterable):
699 "Returns True if all the elements are equal to each other"
700 g = groupby(iterable)
701 return next(g, True) and not next(g, False)
702
Raymond Hettinger30c73622010-12-01 23:45:20 +0000703 def quantify(iterable, pred=bool):
704 "Count how many times the predicate is true"
705 return sum(map(pred, iterable))
Georg Brandl116aa622007-08-15 14:28:22 +0000706
Raymond Hettinger30c73622010-12-01 23:45:20 +0000707 def padnone(iterable):
708 """Returns the sequence elements and then returns None indefinitely.
Georg Brandl116aa622007-08-15 14:28:22 +0000709
Raymond Hettinger30c73622010-12-01 23:45:20 +0000710 Useful for emulating the behavior of the built-in map() function.
711 """
712 return chain(iterable, repeat(None))
Georg Brandl116aa622007-08-15 14:28:22 +0000713
Raymond Hettinger30c73622010-12-01 23:45:20 +0000714 def ncycles(iterable, n):
715 "Returns the sequence elements n times"
716 return chain.from_iterable(repeat(tuple(iterable), n))
Georg Brandl116aa622007-08-15 14:28:22 +0000717
Raymond Hettinger30c73622010-12-01 23:45:20 +0000718 def dotproduct(vec1, vec2):
719 return sum(map(operator.mul, vec1, vec2))
Georg Brandl116aa622007-08-15 14:28:22 +0000720
Raymond Hettinger30c73622010-12-01 23:45:20 +0000721 def flatten(listOfLists):
722 "Flatten one level of nesting"
723 return chain.from_iterable(listOfLists)
Georg Brandl116aa622007-08-15 14:28:22 +0000724
Raymond Hettinger30c73622010-12-01 23:45:20 +0000725 def repeatfunc(func, times=None, *args):
726 """Repeat calls to func with specified arguments.
Georg Brandl116aa622007-08-15 14:28:22 +0000727
Raymond Hettinger30c73622010-12-01 23:45:20 +0000728 Example: repeatfunc(random.random)
729 """
730 if times is None:
731 return starmap(func, repeat(args))
732 return starmap(func, repeat(args, times))
Georg Brandl116aa622007-08-15 14:28:22 +0000733
Raymond Hettinger30c73622010-12-01 23:45:20 +0000734 def pairwise(iterable):
735 "s -> (s0,s1), (s1,s2), (s2, s3), ..."
736 a, b = tee(iterable)
737 next(b, None)
738 return zip(a, b)
Georg Brandl116aa622007-08-15 14:28:22 +0000739
Raymond Hettinger44571da2013-05-05 19:53:41 -0700740 def grouper(iterable, n, fillvalue=None):
Raymond Hettinger9ae94732012-07-08 16:04:03 -0700741 "Collect data into fixed-length chunks or blocks"
Raymond Hettinger44571da2013-05-05 19:53:41 -0700742 # grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx"
Raymond Hettinger30c73622010-12-01 23:45:20 +0000743 args = [iter(iterable)] * n
744 return zip_longest(*args, fillvalue=fillvalue)
Georg Brandl116aa622007-08-15 14:28:22 +0000745
Raymond Hettinger30c73622010-12-01 23:45:20 +0000746 def roundrobin(*iterables):
747 "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
748 # Recipe credited to George Sakkis
749 pending = len(iterables)
750 nexts = cycle(iter(it).__next__ for it in iterables)
751 while pending:
752 try:
753 for next in nexts:
754 yield next()
755 except StopIteration:
756 pending -= 1
757 nexts = cycle(islice(nexts, pending))
Georg Brandl116aa622007-08-15 14:28:22 +0000758
Raymond Hettinger30c73622010-12-01 23:45:20 +0000759 def partition(pred, iterable):
760 'Use a predicate to partition entries into false entries and true entries'
761 # partition(is_odd, range(10)) --> 0 2 4 6 8 and 1 3 5 7 9
762 t1, t2 = tee(iterable)
763 return filterfalse(pred, t1), filter(pred, t2)
Raymond Hettinger5ce0aa22010-12-01 10:49:19 +0000764
Raymond Hettinger30c73622010-12-01 23:45:20 +0000765 def powerset(iterable):
766 "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
767 s = list(iterable)
768 return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettinger08d01ee2010-08-07 05:36:53 +0000769
Raymond Hettinger30c73622010-12-01 23:45:20 +0000770 def unique_everseen(iterable, key=None):
771 "List unique elements, preserving order. Remember all elements ever seen."
772 # unique_everseen('AAAABBBCCDAABBB') --> A B C D
773 # unique_everseen('ABBCcAD', str.lower) --> A B C D
774 seen = set()
775 seen_add = seen.add
776 if key is None:
777 for element in filterfalse(seen.__contains__, iterable):
778 seen_add(element)
779 yield element
780 else:
781 for element in iterable:
782 k = key(element)
783 if k not in seen:
784 seen_add(k)
785 yield element
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000786
Raymond Hettinger30c73622010-12-01 23:45:20 +0000787 def unique_justseen(iterable, key=None):
788 "List unique elements, preserving order. Remember only the element just seen."
789 # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
790 # unique_justseen('ABBCcAD', str.lower) --> A B C A D
791 return map(next, map(itemgetter(1), groupby(iterable, key)))
Raymond Hettingerad9d96b2009-01-02 21:39:07 +0000792
Raymond Hettinger30c73622010-12-01 23:45:20 +0000793 def iter_except(func, exception, first=None):
794 """ Call a function repeatedly until an exception is raised.
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000795
Raymond Hettinger30c73622010-12-01 23:45:20 +0000796 Converts a call-until-exception interface to an iterator interface.
Andrew Kuchling1d7d5802013-06-20 21:17:41 -0400797 Like builtins.iter(func, sentinel) but uses an exception instead
Raymond Hettinger30c73622010-12-01 23:45:20 +0000798 of a sentinel to end the loop.
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000799
Raymond Hettinger30c73622010-12-01 23:45:20 +0000800 Examples:
801 iter_except(functools.partial(heappop, h), IndexError) # priority queue iterator
802 iter_except(d.popitem, KeyError) # non-blocking dict iterator
803 iter_except(d.popleft, IndexError) # non-blocking deque iterator
804 iter_except(q.get_nowait, Queue.Empty) # loop over a producer Queue
805 iter_except(s.pop, KeyError) # non-blocking set iterator
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000806
Raymond Hettinger30c73622010-12-01 23:45:20 +0000807 """
808 try:
809 if first is not None:
810 yield first() # For database APIs needing an initial cast to db.first()
Raymond Hettingera503f702016-03-13 00:12:31 -0800811 while True:
Raymond Hettinger30c73622010-12-01 23:45:20 +0000812 yield func()
813 except exception:
814 pass
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000815
Raymond Hettinger31b26f62014-04-02 03:16:42 -0700816 def first_true(iterable, default=False, pred=None):
817 """Returns the first true value in the iterable.
818
819 If no true value is found, returns *default*
820
821 If *pred* is not None, returns the first item
822 for which pred(item) is true.
823
824 """
825 # first_true([a,b,c], x) --> a or b or c or x
826 # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
827 return next(filter(pred, iterable), default)
828
Raymond Hettinger30c73622010-12-01 23:45:20 +0000829 def random_product(*args, repeat=1):
830 "Random selection from itertools.product(*args, **kwds)"
831 pools = [tuple(pool) for pool in args] * repeat
832 return tuple(random.choice(pool) for pool in pools)
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000833
Raymond Hettinger30c73622010-12-01 23:45:20 +0000834 def random_permutation(iterable, r=None):
835 "Random selection from itertools.permutations(iterable, r)"
836 pool = tuple(iterable)
837 r = len(pool) if r is None else r
838 return tuple(random.sample(pool, r))
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000839
Raymond Hettinger30c73622010-12-01 23:45:20 +0000840 def random_combination(iterable, r):
841 "Random selection from itertools.combinations(iterable, r)"
842 pool = tuple(iterable)
843 n = len(pool)
844 indices = sorted(random.sample(range(n), r))
845 return tuple(pool[i] for i in indices)
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000846
Raymond Hettinger30c73622010-12-01 23:45:20 +0000847 def random_combination_with_replacement(iterable, r):
848 "Random selection from itertools.combinations_with_replacement(iterable, r)"
849 pool = tuple(iterable)
850 n = len(pool)
851 indices = sorted(random.randrange(n) for i in range(r))
852 return tuple(pool[i] for i in indices)
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000853
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000854Note, many of the above recipes can be optimized by replacing global lookups
855with local variables defined as default values. For example, the
856*dotproduct* recipe can be written as::
857
Raymond Hettinger30c73622010-12-01 23:45:20 +0000858 def dotproduct(vec1, vec2, sum=sum, map=map, mul=operator.mul):
859 return sum(map(mul, vec1, vec2))