blob: 107bc515a677856ac957a4e429d8447c7f80ab15 [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============================================== ==================== =============================================================
Georg Brandl116aa622007-08-15 14:28:22 +000074
Benedikt Werner14e3c442019-03-21 16:28:49 +010075============================================== =============================================================
76Examples Results
77============================================== =============================================================
78``product('ABCD', repeat=2)`` ``AA AB AC AD BA BB BC BD CA CB CC CD DA DB DC DD``
79``permutations('ABCD', 2)`` ``AB AC AD BA BC BD CA CB CD DA DB DC``
80``combinations('ABCD', 2)`` ``AB AC AD BC BD CD``
81``combinations_with_replacement('ABCD'2)`` ``AA AB AC AD BB BC BD CC CD DD``
82============================================== =============================================================
83
Georg Brandl116aa622007-08-15 14:28:22 +000084
85.. _itertools-functions:
86
87Itertool functions
88------------------
89
90The following module functions all construct and return iterators. Some provide
91streams of infinite length, so they should only be accessed by functions or
92loops that truncate the stream.
93
Lisa Roach9718b592018-09-23 17:34:59 -070094.. function:: accumulate(iterable[, func, *, initial=None])
Raymond Hettingeradb81462010-12-01 22:50:36 +000095
Andrew Kuchling15b04eb2014-04-15 22:28:40 -040096 Make an iterator that returns accumulated sums, or accumulated
97 results of other binary functions (specified via the optional
Lisa Roach9718b592018-09-23 17:34:59 -070098 *func* argument).
99
100 If *func* is supplied, it should be a function
Andrew Kuchling15b04eb2014-04-15 22:28:40 -0400101 of two arguments. Elements of the input *iterable* may be any type
102 that can be accepted as arguments to *func*. (For example, with
103 the default operation of addition, elements may be any addable
104 type including :class:`~decimal.Decimal` or
Lisa Roach9718b592018-09-23 17:34:59 -0700105 :class:`~fractions.Fraction`.)
106
107 Usually, the number of elements output matches the input iterable.
108 However, if the keyword argument *initial* is provided, the
109 accumulation leads off with the *initial* value so that the output
110 has one more element than the input iterable.
Raymond Hettingeradb81462010-12-01 22:50:36 +0000111
Raymond Hettinger672866d2016-05-28 00:17:54 -0700112 Roughly equivalent to::
Raymond Hettinger5d446132011-03-27 18:52:10 -0700113
Lisa Roach9718b592018-09-23 17:34:59 -0700114 def accumulate(iterable, func=operator.add, *, initial=None):
Raymond Hettingeradb81462010-12-01 22:50:36 +0000115 'Return running totals'
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000116 # accumulate([1,2,3,4,5]) --> 1 3 6 10 15
Lisa Roach9718b592018-09-23 17:34:59 -0700117 # accumulate([1,2,3,4,5], initial=100) --> 100 101 103 106 110 115
Raymond Hettinger5d446132011-03-27 18:52:10 -0700118 # accumulate([1,2,3,4,5], operator.mul) --> 1 2 6 24 120
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000119 it = iter(iterable)
Lisa Roach9718b592018-09-23 17:34:59 -0700120 total = initial
121 if initial is None:
122 try:
123 total = next(it)
124 except StopIteration:
125 return
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000126 yield total
127 for element in it:
Raymond Hettinger5d446132011-03-27 18:52:10 -0700128 total = func(total, element)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000129 yield total
Raymond Hettingeradb81462010-12-01 22:50:36 +0000130
Raymond Hettinger295c1d42011-04-21 11:09:28 -0700131 There are a number of uses for the *func* argument. It can be set to
132 :func:`min` for a running minimum, :func:`max` for a running maximum, or
133 :func:`operator.mul` for a running product. Amortization tables can be
134 built by accumulating interest and applying payments. First-order
Georg Brandl5d941342016-02-26 19:37:12 +0100135 `recurrence relations <https://en.wikipedia.org/wiki/Recurrence_relation>`_
Raymond Hettinger295c1d42011-04-21 11:09:28 -0700136 can be modeled by supplying the initial value in the iterable and using only
137 the accumulated total in *func* argument::
Raymond Hettinger5d446132011-03-27 18:52:10 -0700138
139 >>> data = [3, 4, 6, 2, 1, 9, 0, 7, 5, 8]
140 >>> list(accumulate(data, operator.mul)) # running product
141 [3, 12, 72, 144, 144, 1296, 0, 0, 0, 0]
142 >>> list(accumulate(data, max)) # running maximum
143 [3, 4, 6, 6, 6, 9, 9, 9, 9, 9]
144
145 # Amortize a 5% loan of 1000 with 4 annual payments of 90
146 >>> cashflows = [1000, -90, -90, -90, -90]
147 >>> list(accumulate(cashflows, lambda bal, pmt: bal*1.05 + pmt))
148 [1000, 960.0, 918.0, 873.9000000000001, 827.5950000000001]
149
Georg Brandl5d941342016-02-26 19:37:12 +0100150 # Chaotic recurrence relation https://en.wikipedia.org/wiki/Logistic_map
Raymond Hettinger295c1d42011-04-21 11:09:28 -0700151 >>> logistic_map = lambda x, _: r * x * (1 - x)
152 >>> r = 3.8
153 >>> x0 = 0.4
154 >>> inputs = repeat(x0, 36) # only the initial value is used
155 >>> [format(x, '.2f') for x in accumulate(inputs, logistic_map)]
156 ['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 +0200157 '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 -0700158 '0.93', '0.25', '0.71', '0.79', '0.63', '0.88', '0.39', '0.91', '0.32',
159 '0.83', '0.54', '0.95', '0.20', '0.60', '0.91', '0.30', '0.80', '0.60']
160
Raymond Hettinger64801682013-10-12 16:04:17 -0700161 See :func:`functools.reduce` for a similar function that returns only the
162 final accumulated value.
163
Raymond Hettingeradb81462010-12-01 22:50:36 +0000164 .. versionadded:: 3.2
Georg Brandl116aa622007-08-15 14:28:22 +0000165
Raymond Hettinger5d446132011-03-27 18:52:10 -0700166 .. versionchanged:: 3.3
167 Added the optional *func* parameter.
168
Lisa Roach9718b592018-09-23 17:34:59 -0700169 .. versionchanged:: 3.8
170 Added the optional *initial* parameter.
171
Georg Brandl116aa622007-08-15 14:28:22 +0000172.. function:: chain(*iterables)
173
Raymond Hettinger30c73622010-12-01 23:45:20 +0000174 Make an iterator that returns elements from the first iterable until it is
175 exhausted, then proceeds to the next iterable, until all of the iterables are
176 exhausted. Used for treating consecutive sequences as a single sequence.
Raymond Hettinger672866d2016-05-28 00:17:54 -0700177 Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000178
Raymond Hettinger30c73622010-12-01 23:45:20 +0000179 def chain(*iterables):
180 # chain('ABC', 'DEF') --> A B C D E F
181 for it in iterables:
182 for element in it:
183 yield element
Georg Brandl116aa622007-08-15 14:28:22 +0000184
185
Georg Brandl933b9742010-07-29 14:36:11 +0000186.. classmethod:: chain.from_iterable(iterable)
Christian Heimes70e7ea22008-02-28 20:02:27 +0000187
Raymond Hettinger30c73622010-12-01 23:45:20 +0000188 Alternate constructor for :func:`chain`. Gets chained inputs from a
Raymond Hettinger1e21ebc2013-09-09 01:54:27 -0500189 single iterable argument that is evaluated lazily. Roughly equivalent to::
Christian Heimes70e7ea22008-02-28 20:02:27 +0000190
Raymond Hettinger30c73622010-12-01 23:45:20 +0000191 def from_iterable(iterables):
192 # chain.from_iterable(['ABC', 'DEF']) --> A B C D E F
193 for it in iterables:
194 for element in it:
195 yield element
Christian Heimes70e7ea22008-02-28 20:02:27 +0000196
Christian Heimes78644762008-03-04 23:39:23 +0000197
Christian Heimes836baa52008-02-26 08:18:30 +0000198.. function:: combinations(iterable, r)
199
Raymond Hettinger30c73622010-12-01 23:45:20 +0000200 Return *r* length subsequences of elements from the input *iterable*.
Christian Heimes836baa52008-02-26 08:18:30 +0000201
Ruaridh Williamson5e0ed8a2020-05-28 20:56:43 +0100202 The combination tuples are emitted in lexicographic ordering according to
203 the order of the input *iterable*. So, if the input *iterable* is sorted,
204 the combination tuples will be produced in sorted order.
Christian Heimes836baa52008-02-26 08:18:30 +0000205
Raymond Hettinger30c73622010-12-01 23:45:20 +0000206 Elements are treated as unique based on their position, not on their
207 value. So if the input elements are unique, there will be no repeat
208 values in each combination.
Christian Heimes836baa52008-02-26 08:18:30 +0000209
Raymond Hettinger672866d2016-05-28 00:17:54 -0700210 Roughly equivalent to::
Christian Heimes836baa52008-02-26 08:18:30 +0000211
212 def combinations(iterable, r):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000213 # combinations('ABCD', 2) --> AB AC AD BC BD CD
214 # combinations(range(4), 3) --> 012 013 023 123
Christian Heimes836baa52008-02-26 08:18:30 +0000215 pool = tuple(iterable)
Christian Heimes380f7f22008-02-28 11:19:05 +0000216 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000217 if r > n:
218 return
219 indices = list(range(r))
Christian Heimesb558a2e2008-03-02 22:46:37 +0000220 yield tuple(pool[i] for i in indices)
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000221 while True:
Christian Heimes380f7f22008-02-28 11:19:05 +0000222 for i in reversed(range(r)):
Christian Heimesb558a2e2008-03-02 22:46:37 +0000223 if indices[i] != i + n - r:
Christian Heimes836baa52008-02-26 08:18:30 +0000224 break
Christian Heimes380f7f22008-02-28 11:19:05 +0000225 else:
226 return
Christian Heimesb558a2e2008-03-02 22:46:37 +0000227 indices[i] += 1
Christian Heimes380f7f22008-02-28 11:19:05 +0000228 for j in range(i+1, r):
Christian Heimesb558a2e2008-03-02 22:46:37 +0000229 indices[j] = indices[j-1] + 1
230 yield tuple(pool[i] for i in indices)
Christian Heimes836baa52008-02-26 08:18:30 +0000231
Raymond Hettinger30c73622010-12-01 23:45:20 +0000232 The code for :func:`combinations` can be also expressed as a subsequence
233 of :func:`permutations` after filtering entries where the elements are not
234 in sorted order (according to their position in the input pool)::
Christian Heimes78644762008-03-04 23:39:23 +0000235
236 def combinations(iterable, r):
237 pool = tuple(iterable)
238 n = len(pool)
239 for indices in permutations(range(n), r):
240 if sorted(indices) == list(indices):
241 yield tuple(pool[i] for i in indices)
242
Raymond Hettinger30c73622010-12-01 23:45:20 +0000243 The number of items returned is ``n! / r! / (n-r)!`` when ``0 <= r <= n``
244 or zero when ``r > n``.
Christian Heimes836baa52008-02-26 08:18:30 +0000245
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000246.. function:: combinations_with_replacement(iterable, r)
247
Raymond Hettinger30c73622010-12-01 23:45:20 +0000248 Return *r* length subsequences of elements from the input *iterable*
249 allowing individual elements to be repeated more than once.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000250
Ruaridh Williamson5e0ed8a2020-05-28 20:56:43 +0100251 The combination tuples are emitted in lexicographic ordering according to
252 the order of the input *iterable*. So, if the input *iterable* is sorted,
253 the combination tuples will be produced in sorted order.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000254
Raymond Hettinger30c73622010-12-01 23:45:20 +0000255 Elements are treated as unique based on their position, not on their
256 value. So if the input elements are unique, the generated combinations
257 will also be unique.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000258
Raymond Hettinger672866d2016-05-28 00:17:54 -0700259 Roughly equivalent to::
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000260
261 def combinations_with_replacement(iterable, r):
262 # combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC
263 pool = tuple(iterable)
264 n = len(pool)
265 if not n and r:
266 return
267 indices = [0] * r
268 yield tuple(pool[i] for i in indices)
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000269 while True:
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000270 for i in reversed(range(r)):
271 if indices[i] != n - 1:
272 break
273 else:
274 return
275 indices[i:] = [indices[i] + 1] * (r - i)
276 yield tuple(pool[i] for i in indices)
277
Raymond Hettinger30c73622010-12-01 23:45:20 +0000278 The code for :func:`combinations_with_replacement` can be also expressed as
279 a subsequence of :func:`product` after filtering entries where the elements
280 are not in sorted order (according to their position in the input pool)::
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000281
282 def combinations_with_replacement(iterable, r):
283 pool = tuple(iterable)
284 n = len(pool)
285 for indices in product(range(n), repeat=r):
286 if sorted(indices) == list(indices):
287 yield tuple(pool[i] for i in indices)
288
Raymond Hettinger30c73622010-12-01 23:45:20 +0000289 The number of items returned is ``(n+r-1)! / r! / (n-1)!`` when ``n > 0``.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000290
Raymond Hettinger30c73622010-12-01 23:45:20 +0000291 .. versionadded:: 3.1
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000292
Georg Brandl67b21b72010-08-17 15:07:14 +0000293
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000294.. function:: compress(data, selectors)
295
Raymond Hettinger30c73622010-12-01 23:45:20 +0000296 Make an iterator that filters elements from *data* returning only those that
297 have a corresponding element in *selectors* that evaluates to ``True``.
298 Stops when either the *data* or *selectors* iterables has been exhausted.
Raymond Hettinger672866d2016-05-28 00:17:54 -0700299 Roughly equivalent to::
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000300
Raymond Hettinger30c73622010-12-01 23:45:20 +0000301 def compress(data, selectors):
302 # compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F
303 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000304
Raymond Hettinger30c73622010-12-01 23:45:20 +0000305 .. versionadded:: 3.1
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000306
307
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000308.. function:: count(start=0, step=1)
Georg Brandl116aa622007-08-15 14:28:22 +0000309
Andrew Kuchlingedb42602013-06-21 08:00:58 -0400310 Make an iterator that returns evenly spaced values starting with number *start*. Often
Raymond Hettinger30c73622010-12-01 23:45:20 +0000311 used as an argument to :func:`map` to generate consecutive data points.
Raymond Hettinger672866d2016-05-28 00:17:54 -0700312 Also, used with :func:`zip` to add sequence numbers. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000313
Raymond Hettinger30c73622010-12-01 23:45:20 +0000314 def count(start=0, step=1):
315 # count(10) --> 10 11 12 13 14 ...
Georg Brandl37a80dc2011-01-13 07:31:18 +0000316 # count(2.5, 0.5) -> 2.5 3.0 3.5 ...
Raymond Hettinger30c73622010-12-01 23:45:20 +0000317 n = start
318 while True:
319 yield n
320 n += step
Georg Brandl116aa622007-08-15 14:28:22 +0000321
Raymond Hettinger30c73622010-12-01 23:45:20 +0000322 When counting with floating point numbers, better accuracy can sometimes be
323 achieved by substituting multiplicative code such as: ``(start + step * i
324 for i in count())``.
Raymond Hettinger5bc472a2009-06-17 01:40:52 +0000325
Raymond Hettinger30c73622010-12-01 23:45:20 +0000326 .. versionchanged:: 3.1
327 Added *step* argument and allowed non-integer arguments.
Georg Brandl116aa622007-08-15 14:28:22 +0000328
329.. function:: cycle(iterable)
330
Raymond Hettinger30c73622010-12-01 23:45:20 +0000331 Make an iterator returning elements from the iterable and saving a copy of each.
332 When the iterable is exhausted, return elements from the saved copy. Repeats
Raymond Hettinger672866d2016-05-28 00:17:54 -0700333 indefinitely. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000334
Raymond Hettinger30c73622010-12-01 23:45:20 +0000335 def cycle(iterable):
336 # cycle('ABCD') --> A B C D A B C D A B C D ...
337 saved = []
338 for element in iterable:
339 yield element
340 saved.append(element)
341 while saved:
342 for element in saved:
Georg Brandl116aa622007-08-15 14:28:22 +0000343 yield element
344
Raymond Hettinger30c73622010-12-01 23:45:20 +0000345 Note, this member of the toolkit may require significant auxiliary storage
346 (depending on the length of the iterable).
Georg Brandl116aa622007-08-15 14:28:22 +0000347
348
349.. function:: dropwhile(predicate, iterable)
350
Raymond Hettinger30c73622010-12-01 23:45:20 +0000351 Make an iterator that drops elements from the iterable as long as the predicate
352 is true; afterwards, returns every element. Note, the iterator does not produce
353 *any* output until the predicate first becomes false, so it may have a lengthy
Raymond Hettinger672866d2016-05-28 00:17:54 -0700354 start-up time. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000355
Raymond Hettinger30c73622010-12-01 23:45:20 +0000356 def dropwhile(predicate, iterable):
357 # dropwhile(lambda x: x<5, [1,4,6,4,1]) --> 6 4 1
358 iterable = iter(iterable)
359 for x in iterable:
360 if not predicate(x):
361 yield x
362 break
363 for x in iterable:
364 yield x
Georg Brandl116aa622007-08-15 14:28:22 +0000365
Raymond Hettinger749761e2009-01-27 04:42:48 +0000366.. function:: filterfalse(predicate, iterable)
367
Raymond Hettinger30c73622010-12-01 23:45:20 +0000368 Make an iterator that filters elements from iterable returning only those for
369 which the predicate is ``False``. If *predicate* is ``None``, return the items
Raymond Hettinger672866d2016-05-28 00:17:54 -0700370 that are false. Roughly equivalent to::
Raymond Hettinger749761e2009-01-27 04:42:48 +0000371
Raymond Hettinger30c73622010-12-01 23:45:20 +0000372 def filterfalse(predicate, iterable):
373 # filterfalse(lambda x: x%2, range(10)) --> 0 2 4 6 8
374 if predicate is None:
375 predicate = bool
376 for x in iterable:
377 if not predicate(x):
378 yield x
Raymond Hettinger749761e2009-01-27 04:42:48 +0000379
Georg Brandl116aa622007-08-15 14:28:22 +0000380
Georg Brandl3dd33882009-06-01 17:35:27 +0000381.. function:: groupby(iterable, key=None)
Georg Brandl116aa622007-08-15 14:28:22 +0000382
Raymond Hettinger30c73622010-12-01 23:45:20 +0000383 Make an iterator that returns consecutive keys and groups from the *iterable*.
384 The *key* is a function computing a key value for each element. If not
385 specified or is ``None``, *key* defaults to an identity function and returns
386 the element unchanged. Generally, the iterable needs to already be sorted on
387 the same key function.
Georg Brandl116aa622007-08-15 14:28:22 +0000388
Raymond Hettinger30c73622010-12-01 23:45:20 +0000389 The operation of :func:`groupby` is similar to the ``uniq`` filter in Unix. It
390 generates a break or new group every time the value of the key function changes
391 (which is why it is usually necessary to have sorted the data using the same key
392 function). That behavior differs from SQL's GROUP BY which aggregates common
393 elements regardless of their input order.
Georg Brandl116aa622007-08-15 14:28:22 +0000394
Raymond Hettinger30c73622010-12-01 23:45:20 +0000395 The returned group is itself an iterator that shares the underlying iterable
396 with :func:`groupby`. Because the source is shared, when the :func:`groupby`
397 object is advanced, the previous group is no longer visible. So, if that data
398 is needed later, it should be stored as a list::
Georg Brandl116aa622007-08-15 14:28:22 +0000399
Raymond Hettinger30c73622010-12-01 23:45:20 +0000400 groups = []
401 uniquekeys = []
402 data = sorted(data, key=keyfunc)
403 for k, g in groupby(data, keyfunc):
404 groups.append(list(g)) # Store group iterator as a list
405 uniquekeys.append(k)
Georg Brandl116aa622007-08-15 14:28:22 +0000406
Raymond Hettinger672866d2016-05-28 00:17:54 -0700407 :func:`groupby` is roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000408
Raymond Hettinger30c73622010-12-01 23:45:20 +0000409 class groupby:
410 # [k for k, g in groupby('AAAABBBCCDAABBB')] --> A B C D A B
411 # [list(g) for k, g in groupby('AAAABBBCCD')] --> AAAA BBB CC D
412 def __init__(self, iterable, key=None):
413 if key is None:
414 key = lambda x: x
415 self.keyfunc = key
416 self.it = iter(iterable)
417 self.tgtkey = self.currkey = self.currvalue = object()
418 def __iter__(self):
419 return self
420 def __next__(self):
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300421 self.id = object()
Raymond Hettinger30c73622010-12-01 23:45:20 +0000422 while self.currkey == self.tgtkey:
423 self.currvalue = next(self.it) # Exit on StopIteration
424 self.currkey = self.keyfunc(self.currvalue)
425 self.tgtkey = self.currkey
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300426 return (self.currkey, self._grouper(self.tgtkey, self.id))
427 def _grouper(self, tgtkey, id):
428 while self.id is id and self.currkey == tgtkey:
Raymond Hettinger30c73622010-12-01 23:45:20 +0000429 yield self.currvalue
Raymond Hettinger828d9322014-11-22 21:56:23 -0800430 try:
431 self.currvalue = next(self.it)
432 except StopIteration:
433 return
Raymond Hettinger30c73622010-12-01 23:45:20 +0000434 self.currkey = self.keyfunc(self.currvalue)
Georg Brandl116aa622007-08-15 14:28:22 +0000435
Georg Brandl116aa622007-08-15 14:28:22 +0000436
Ezio Melottie0add762012-09-14 06:32:35 +0300437.. function:: islice(iterable, stop)
438 islice(iterable, start, stop[, step])
Georg Brandl116aa622007-08-15 14:28:22 +0000439
Raymond Hettinger30c73622010-12-01 23:45:20 +0000440 Make an iterator that returns selected elements from the iterable. If *start* is
441 non-zero, then elements from the iterable are skipped until start is reached.
442 Afterward, elements are returned consecutively unless *step* is set higher than
443 one which results in items being skipped. If *stop* is ``None``, then iteration
444 continues until the iterator is exhausted, if at all; otherwise, it stops at the
445 specified position. Unlike regular slicing, :func:`islice` does not support
446 negative values for *start*, *stop*, or *step*. Can be used to extract related
447 fields from data where the internal structure has been flattened (for example, a
Raymond Hettinger672866d2016-05-28 00:17:54 -0700448 multi-line report may list a name field on every third line). Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000449
Raymond Hettinger30c73622010-12-01 23:45:20 +0000450 def islice(iterable, *args):
451 # islice('ABCDEFG', 2) --> A B
452 # islice('ABCDEFG', 2, 4) --> C D
453 # islice('ABCDEFG', 2, None) --> C D E F G
454 # islice('ABCDEFG', 0, None, 2) --> A C E G
455 s = slice(*args)
Cheryl Sabellada1734c2018-03-26 21:29:33 -0400456 start, stop, step = s.start or 0, s.stop or sys.maxsize, s.step or 1
457 it = iter(range(start, stop, step))
Raymond Hettinger828d9322014-11-22 21:56:23 -0800458 try:
459 nexti = next(it)
460 except StopIteration:
Cheryl Sabellada1734c2018-03-26 21:29:33 -0400461 # Consume *iterable* up to the *start* position.
462 for i, element in zip(range(start), iterable):
463 pass
Raymond Hettinger828d9322014-11-22 21:56:23 -0800464 return
Cheryl Sabellada1734c2018-03-26 21:29:33 -0400465 try:
466 for i, element in enumerate(iterable):
467 if i == nexti:
468 yield element
469 nexti = next(it)
470 except StopIteration:
471 # Consume to *stop*.
472 for i, element in zip(range(i + 1, stop), iterable):
473 pass
Georg Brandl116aa622007-08-15 14:28:22 +0000474
Raymond Hettinger30c73622010-12-01 23:45:20 +0000475 If *start* is ``None``, then iteration starts at zero. If *step* is ``None``,
476 then the step defaults to one.
Georg Brandl116aa622007-08-15 14:28:22 +0000477
Georg Brandl116aa622007-08-15 14:28:22 +0000478
Georg Brandl3dd33882009-06-01 17:35:27 +0000479.. function:: permutations(iterable, r=None)
Christian Heimes70e7ea22008-02-28 20:02:27 +0000480
Raymond Hettinger30c73622010-12-01 23:45:20 +0000481 Return successive *r* length permutations of elements in the *iterable*.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000482
Raymond Hettinger30c73622010-12-01 23:45:20 +0000483 If *r* is not specified or is ``None``, then *r* defaults to the length
484 of the *iterable* and all possible full-length permutations
485 are generated.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000486
Ruaridh Williamson5e0ed8a2020-05-28 20:56:43 +0100487 The permutation tuples are emitted in lexicographic ordering according to
488 the order of the input *iterable*. So, if the input *iterable* is sorted,
489 the combination tuples will be produced in sorted order.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000490
Raymond Hettinger30c73622010-12-01 23:45:20 +0000491 Elements are treated as unique based on their position, not on their
492 value. So if the input elements are unique, there will be no repeat
493 values in each permutation.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000494
Raymond Hettinger672866d2016-05-28 00:17:54 -0700495 Roughly equivalent to::
Christian Heimesb558a2e2008-03-02 22:46:37 +0000496
497 def permutations(iterable, r=None):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000498 # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
499 # permutations(range(3)) --> 012 021 102 120 201 210
Christian Heimesb558a2e2008-03-02 22:46:37 +0000500 pool = tuple(iterable)
501 n = len(pool)
502 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000503 if r > n:
504 return
505 indices = list(range(n))
Sandro Tosi73866622011-12-25 17:25:45 +0100506 cycles = list(range(n, n-r, -1))
Christian Heimesb558a2e2008-03-02 22:46:37 +0000507 yield tuple(pool[i] for i in indices[:r])
508 while n:
509 for i in reversed(range(r)):
510 cycles[i] -= 1
511 if cycles[i] == 0:
512 indices[i:] = indices[i+1:] + indices[i:i+1]
513 cycles[i] = n - i
514 else:
515 j = cycles[i]
516 indices[i], indices[-j] = indices[-j], indices[i]
517 yield tuple(pool[i] for i in indices[:r])
518 break
519 else:
520 return
Christian Heimes70e7ea22008-02-28 20:02:27 +0000521
Raymond Hettinger30c73622010-12-01 23:45:20 +0000522 The code for :func:`permutations` can be also expressed as a subsequence of
523 :func:`product`, filtered to exclude entries with repeated elements (those
524 from the same position in the input pool)::
Christian Heimes78644762008-03-04 23:39:23 +0000525
526 def permutations(iterable, r=None):
527 pool = tuple(iterable)
528 n = len(pool)
529 r = n if r is None else r
530 for indices in product(range(n), repeat=r):
531 if len(set(indices)) == r:
532 yield tuple(pool[i] for i in indices)
533
Raymond Hettinger30c73622010-12-01 23:45:20 +0000534 The number of items returned is ``n! / (n-r)!`` when ``0 <= r <= n``
535 or zero when ``r > n``.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000536
Georg Brandl3dd33882009-06-01 17:35:27 +0000537.. function:: product(*iterables, repeat=1)
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000538
Raymond Hettinger30c73622010-12-01 23:45:20 +0000539 Cartesian product of input iterables.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000540
Raymond Hettinger672866d2016-05-28 00:17:54 -0700541 Roughly equivalent to nested for-loops in a generator expression. For example,
Raymond Hettinger30c73622010-12-01 23:45:20 +0000542 ``product(A, B)`` returns the same as ``((x,y) for x in A for y in B)``.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000543
Raymond Hettinger30c73622010-12-01 23:45:20 +0000544 The nested loops cycle like an odometer with the rightmost element advancing
545 on every iteration. This pattern creates a lexicographic ordering so that if
546 the input's iterables are sorted, the product tuples are emitted in sorted
547 order.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000548
Raymond Hettinger30c73622010-12-01 23:45:20 +0000549 To compute the product of an iterable with itself, specify the number of
550 repetitions with the optional *repeat* keyword argument. For example,
551 ``product(A, repeat=4)`` means the same as ``product(A, A, A, A)``.
Christian Heimes9e7f1d22008-02-28 12:27:11 +0000552
Raymond Hettinger672866d2016-05-28 00:17:54 -0700553 This function is roughly equivalent to the following code, except that the
Raymond Hettinger30c73622010-12-01 23:45:20 +0000554 actual implementation does not build up intermediate results in memory::
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000555
Raymond Hettinger30c73622010-12-01 23:45:20 +0000556 def product(*args, repeat=1):
557 # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy
558 # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111
559 pools = [tuple(pool) for pool in args] * repeat
560 result = [[]]
561 for pool in pools:
562 result = [x+[y] for x in result for y in pool]
563 for prod in result:
564 yield tuple(prod)
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000565
Ramil Nugmanovcfc6ce42020-05-28 19:46:22 +0300566 Before :func:`product` runs, it completely consumes the input iterables,
567 keeping pools of values in memory to generate the products. Accordingly,
568 it only useful with finite inputs.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000569
Raymond Hettingerd75ad442009-06-01 19:16:52 +0000570.. function:: repeat(object[, times])
Georg Brandl116aa622007-08-15 14:28:22 +0000571
Raymond Hettinger30c73622010-12-01 23:45:20 +0000572 Make an iterator that returns *object* over and over again. Runs indefinitely
573 unless the *times* argument is specified. Used as argument to :func:`map` for
574 invariant parameters to the called function. Also used with :func:`zip` to
Raymond Hettinger672866d2016-05-28 00:17:54 -0700575 create an invariant part of a tuple record.
576
577 Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000578
Raymond Hettinger30c73622010-12-01 23:45:20 +0000579 def repeat(object, times=None):
580 # repeat(10, 3) --> 10 10 10
581 if times is None:
582 while True:
583 yield object
584 else:
585 for i in range(times):
586 yield object
Georg Brandl116aa622007-08-15 14:28:22 +0000587
Raymond Hettingerfc3ba6b2012-02-01 09:07:40 -0800588 A common use for *repeat* is to supply a stream of constant values to *map*
589 or *zip*::
590
591 >>> list(map(pow, range(10), repeat(2)))
592 [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
Georg Brandl116aa622007-08-15 14:28:22 +0000593
594.. function:: starmap(function, iterable)
595
Raymond Hettinger30c73622010-12-01 23:45:20 +0000596 Make an iterator that computes the function using arguments obtained from
597 the iterable. Used instead of :func:`map` when argument parameters are already
598 grouped in tuples from a single iterable (the data has been "pre-zipped"). The
599 difference between :func:`map` and :func:`starmap` parallels the distinction
Raymond Hettinger672866d2016-05-28 00:17:54 -0700600 between ``function(a,b)`` and ``function(*c)``. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000601
Raymond Hettinger30c73622010-12-01 23:45:20 +0000602 def starmap(function, iterable):
603 # starmap(pow, [(2,5), (3,2), (10,3)]) --> 32 9 1000
604 for args in iterable:
605 yield function(*args)
Christian Heimes679db4a2008-01-18 09:56:22 +0000606
Georg Brandl116aa622007-08-15 14:28:22 +0000607
608.. function:: takewhile(predicate, iterable)
609
Raymond Hettinger30c73622010-12-01 23:45:20 +0000610 Make an iterator that returns elements from the iterable as long as the
Raymond Hettinger672866d2016-05-28 00:17:54 -0700611 predicate is true. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000612
Raymond Hettinger30c73622010-12-01 23:45:20 +0000613 def takewhile(predicate, iterable):
614 # takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4
615 for x in iterable:
616 if predicate(x):
617 yield x
618 else:
619 break
Georg Brandl116aa622007-08-15 14:28:22 +0000620
621
Georg Brandl3dd33882009-06-01 17:35:27 +0000622.. function:: tee(iterable, n=2)
Georg Brandl116aa622007-08-15 14:28:22 +0000623
Raymond Hettingerab425aa2016-04-26 00:10:00 -0700624 Return *n* independent iterators from a single iterable.
625
626 The following Python code helps explain what *tee* does (although the actual
Serhiy Storchaka4ecfa452016-05-16 09:31:54 +0300627 implementation is more complex and uses only a single underlying
Raymond Hettinger672866d2016-05-28 00:17:54 -0700628 :abbr:`FIFO (first-in, first-out)` queue).
629
630 Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000631
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000632 def tee(iterable, n=2):
633 it = iter(iterable)
634 deques = [collections.deque() for i in range(n)]
635 def gen(mydeque):
636 while True:
637 if not mydeque: # when the local deque is empty
Raymond Hettinger828d9322014-11-22 21:56:23 -0800638 try:
639 newval = next(it) # fetch a new value and
640 except StopIteration:
641 return
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000642 for d in deques: # load it to all the deques
643 d.append(newval)
644 yield mydeque.popleft()
645 return tuple(gen(d) for d in deques)
Georg Brandl116aa622007-08-15 14:28:22 +0000646
Raymond Hettinger30c73622010-12-01 23:45:20 +0000647 Once :func:`tee` has made a split, the original *iterable* should not be
648 used anywhere else; otherwise, the *iterable* could get advanced without
Serhiy Storchaka918b4682019-09-09 11:18:16 +0300649 the tee objects being informed.
Georg Brandl116aa622007-08-15 14:28:22 +0000650
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300651 ``tee`` iterators are not threadsafe. A :exc:`RuntimeError` may be
652 raised when using simultaneously iterators returned by the same :func:`tee`
653 call, even if the original *iterable* is threadsafe.
654
Raymond Hettinger30c73622010-12-01 23:45:20 +0000655 This itertool may require significant auxiliary storage (depending on how
656 much temporary data needs to be stored). In general, if one iterator uses
657 most or all of the data before another iterator starts, it is faster to use
658 :func:`list` instead of :func:`tee`.
Georg Brandl116aa622007-08-15 14:28:22 +0000659
Georg Brandl116aa622007-08-15 14:28:22 +0000660
Georg Brandl3dd33882009-06-01 17:35:27 +0000661.. function:: zip_longest(*iterables, fillvalue=None)
Raymond Hettinger749761e2009-01-27 04:42:48 +0000662
Raymond Hettinger30c73622010-12-01 23:45:20 +0000663 Make an iterator that aggregates elements from each of the iterables. If the
664 iterables are of uneven length, missing values are filled-in with *fillvalue*.
Raymond Hettinger672866d2016-05-28 00:17:54 -0700665 Iteration continues until the longest iterable is exhausted. Roughly equivalent to::
Raymond Hettinger749761e2009-01-27 04:42:48 +0000666
Raymond Hettinger3147b042017-09-07 14:01:44 -0700667 def zip_longest(*args, fillvalue=None):
Raymond Hettinger30c73622010-12-01 23:45:20 +0000668 # zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-
Raymond Hettinger3147b042017-09-07 14:01:44 -0700669 iterators = [iter(it) for it in args]
670 num_active = len(iterators)
671 if not num_active:
672 return
673 while True:
674 values = []
675 for i, it in enumerate(iterators):
676 try:
677 value = next(it)
678 except StopIteration:
679 num_active -= 1
680 if not num_active:
681 return
682 iterators[i] = repeat(fillvalue)
683 value = fillvalue
684 values.append(value)
685 yield tuple(values)
Raymond Hettinger749761e2009-01-27 04:42:48 +0000686
Raymond Hettinger30c73622010-12-01 23:45:20 +0000687 If one of the iterables is potentially infinite, then the :func:`zip_longest`
688 function should be wrapped with something that limits the number of calls
689 (for example :func:`islice` or :func:`takewhile`). If not specified,
690 *fillvalue* defaults to ``None``.
Raymond Hettinger749761e2009-01-27 04:42:48 +0000691
692
Georg Brandl116aa622007-08-15 14:28:22 +0000693.. _itertools-recipes:
694
Raymond Hettinger1fa76822010-12-06 23:31:36 +0000695Itertools Recipes
696-----------------
Georg Brandl116aa622007-08-15 14:28:22 +0000697
698This section shows recipes for creating an extended toolset using the existing
699itertools as building blocks.
700
Raymond Hettingeradf02b32019-08-04 13:35:58 -0700701Substantially all of these recipes and many, many others can be installed from
702the `more-itertools project <https://pypi.org/project/more-itertools/>`_ found
703on the Python Package Index::
704
705 pip install more-itertools
706
Georg Brandl116aa622007-08-15 14:28:22 +0000707The extended tools offer the same high performance as the underlying toolset.
708The superior memory performance is kept by processing elements one at a time
709rather than bringing the whole iterable into memory all at once. Code volume is
710kept small by linking the tools together in a functional style which helps
711eliminate temporary variables. High speed is retained by preferring
Georg Brandl9afde1c2007-11-01 20:32:30 +0000712"vectorized" building blocks over the use of for-loops and :term:`generator`\s
Christian Heimesfe337bf2008-03-23 21:54:12 +0000713which incur interpreter overhead.
714
715.. testcode::
Georg Brandl116aa622007-08-15 14:28:22 +0000716
Raymond Hettinger30c73622010-12-01 23:45:20 +0000717 def take(n, iterable):
718 "Return first n items of the iterable as a list"
719 return list(islice(iterable, n))
Georg Brandl116aa622007-08-15 14:28:22 +0000720
Raymond Hettinger9265dd72018-04-08 08:44:20 -0700721 def prepend(value, iterator):
722 "Prepend a single value in front of an iterator"
723 # prepend(1, [2, 3, 4]) -> 1 2 3 4
724 return chain([value], iterator)
725
Raymond Hettinger30c73622010-12-01 23:45:20 +0000726 def tabulate(function, start=0):
727 "Return function(0), function(1), ..."
728 return map(function, count(start))
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000729
Raymond Hettingerdfe098d2014-05-25 22:03:56 -0700730 def tail(n, iterable):
731 "Return an iterator over the last n items"
732 # tail(3, 'ABCDEFG') --> E F G
733 return iter(collections.deque(iterable, maxlen=n))
734
Cheryl Sabellada1734c2018-03-26 21:29:33 -0400735 def consume(iterator, n=None):
736 "Advance the iterator n-steps ahead. If n is None, consume entirely."
Raymond Hettinger30c73622010-12-01 23:45:20 +0000737 # Use functions that consume iterators at C speed.
738 if n is None:
739 # feed the entire iterator into a zero-length deque
740 collections.deque(iterator, maxlen=0)
741 else:
742 # advance to the empty slice starting at position n
743 next(islice(iterator, n, n), None)
Raymond Hettingerfa007962009-03-09 11:55:25 +0000744
Raymond Hettinger30c73622010-12-01 23:45:20 +0000745 def nth(iterable, n, default=None):
746 "Returns the nth item or a default value"
747 return next(islice(iterable, n, None), default)
Georg Brandl116aa622007-08-15 14:28:22 +0000748
Raymond Hettingere525ee32016-03-06 18:11:38 -0800749 def all_equal(iterable):
750 "Returns True if all the elements are equal to each other"
751 g = groupby(iterable)
752 return next(g, True) and not next(g, False)
753
Raymond Hettinger30c73622010-12-01 23:45:20 +0000754 def quantify(iterable, pred=bool):
755 "Count how many times the predicate is true"
756 return sum(map(pred, iterable))
Georg Brandl116aa622007-08-15 14:28:22 +0000757
Raymond Hettinger30c73622010-12-01 23:45:20 +0000758 def padnone(iterable):
759 """Returns the sequence elements and then returns None indefinitely.
Georg Brandl116aa622007-08-15 14:28:22 +0000760
Raymond Hettinger30c73622010-12-01 23:45:20 +0000761 Useful for emulating the behavior of the built-in map() function.
762 """
763 return chain(iterable, repeat(None))
Georg Brandl116aa622007-08-15 14:28:22 +0000764
Raymond Hettinger30c73622010-12-01 23:45:20 +0000765 def ncycles(iterable, n):
766 "Returns the sequence elements n times"
767 return chain.from_iterable(repeat(tuple(iterable), n))
Georg Brandl116aa622007-08-15 14:28:22 +0000768
Raymond Hettinger30c73622010-12-01 23:45:20 +0000769 def dotproduct(vec1, vec2):
770 return sum(map(operator.mul, vec1, vec2))
Georg Brandl116aa622007-08-15 14:28:22 +0000771
Борис Верховский99b77012019-11-02 15:09:14 -0400772 def flatten(list_of_lists):
Raymond Hettinger30c73622010-12-01 23:45:20 +0000773 "Flatten one level of nesting"
Борис Верховский99b77012019-11-02 15:09:14 -0400774 return chain.from_iterable(list_of_lists)
Georg Brandl116aa622007-08-15 14:28:22 +0000775
Raymond Hettinger30c73622010-12-01 23:45:20 +0000776 def repeatfunc(func, times=None, *args):
777 """Repeat calls to func with specified arguments.
Georg Brandl116aa622007-08-15 14:28:22 +0000778
Raymond Hettinger30c73622010-12-01 23:45:20 +0000779 Example: repeatfunc(random.random)
780 """
781 if times is None:
782 return starmap(func, repeat(args))
783 return starmap(func, repeat(args, times))
Georg Brandl116aa622007-08-15 14:28:22 +0000784
Raymond Hettinger30c73622010-12-01 23:45:20 +0000785 def pairwise(iterable):
786 "s -> (s0,s1), (s1,s2), (s2, s3), ..."
787 a, b = tee(iterable)
788 next(b, None)
789 return zip(a, b)
Georg Brandl116aa622007-08-15 14:28:22 +0000790
Raymond Hettinger44571da2013-05-05 19:53:41 -0700791 def grouper(iterable, n, fillvalue=None):
Raymond Hettinger9ae94732012-07-08 16:04:03 -0700792 "Collect data into fixed-length chunks or blocks"
Raymond Hettinger44571da2013-05-05 19:53:41 -0700793 # grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx"
Raymond Hettinger30c73622010-12-01 23:45:20 +0000794 args = [iter(iterable)] * n
795 return zip_longest(*args, fillvalue=fillvalue)
Georg Brandl116aa622007-08-15 14:28:22 +0000796
Raymond Hettinger30c73622010-12-01 23:45:20 +0000797 def roundrobin(*iterables):
798 "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
799 # Recipe credited to George Sakkis
Raymond Hettinger337cbba2017-11-21 00:23:34 -0800800 num_active = len(iterables)
Raymond Hettinger30c73622010-12-01 23:45:20 +0000801 nexts = cycle(iter(it).__next__ for it in iterables)
Raymond Hettinger337cbba2017-11-21 00:23:34 -0800802 while num_active:
Raymond Hettinger30c73622010-12-01 23:45:20 +0000803 try:
804 for next in nexts:
805 yield next()
806 except StopIteration:
Raymond Hettinger337cbba2017-11-21 00:23:34 -0800807 # Remove the iterator we just exhausted from the cycle.
808 num_active -= 1
809 nexts = cycle(islice(nexts, num_active))
Georg Brandl116aa622007-08-15 14:28:22 +0000810
Raymond Hettinger30c73622010-12-01 23:45:20 +0000811 def partition(pred, iterable):
812 'Use a predicate to partition entries into false entries and true entries'
813 # partition(is_odd, range(10)) --> 0 2 4 6 8 and 1 3 5 7 9
814 t1, t2 = tee(iterable)
815 return filterfalse(pred, t1), filter(pred, t2)
Raymond Hettinger5ce0aa22010-12-01 10:49:19 +0000816
Raymond Hettinger30c73622010-12-01 23:45:20 +0000817 def powerset(iterable):
818 "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
819 s = list(iterable)
820 return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettinger08d01ee2010-08-07 05:36:53 +0000821
Raymond Hettinger30c73622010-12-01 23:45:20 +0000822 def unique_everseen(iterable, key=None):
823 "List unique elements, preserving order. Remember all elements ever seen."
824 # unique_everseen('AAAABBBCCDAABBB') --> A B C D
825 # unique_everseen('ABBCcAD', str.lower) --> A B C D
826 seen = set()
827 seen_add = seen.add
828 if key is None:
829 for element in filterfalse(seen.__contains__, iterable):
830 seen_add(element)
831 yield element
832 else:
833 for element in iterable:
834 k = key(element)
835 if k not in seen:
836 seen_add(k)
837 yield element
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000838
Raymond Hettinger30c73622010-12-01 23:45:20 +0000839 def unique_justseen(iterable, key=None):
840 "List unique elements, preserving order. Remember only the element just seen."
841 # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
842 # unique_justseen('ABBCcAD', str.lower) --> A B C A D
Jakub Molinskib4c7f392019-04-23 10:30:30 +0200843 return map(next, map(operator.itemgetter(1), groupby(iterable, key)))
Raymond Hettingerad9d96b2009-01-02 21:39:07 +0000844
Raymond Hettinger30c73622010-12-01 23:45:20 +0000845 def iter_except(func, exception, first=None):
846 """ Call a function repeatedly until an exception is raised.
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000847
Raymond Hettinger30c73622010-12-01 23:45:20 +0000848 Converts a call-until-exception interface to an iterator interface.
Andrew Kuchling1d7d5802013-06-20 21:17:41 -0400849 Like builtins.iter(func, sentinel) but uses an exception instead
Raymond Hettinger30c73622010-12-01 23:45:20 +0000850 of a sentinel to end the loop.
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000851
Raymond Hettinger30c73622010-12-01 23:45:20 +0000852 Examples:
853 iter_except(functools.partial(heappop, h), IndexError) # priority queue iterator
854 iter_except(d.popitem, KeyError) # non-blocking dict iterator
855 iter_except(d.popleft, IndexError) # non-blocking deque iterator
856 iter_except(q.get_nowait, Queue.Empty) # loop over a producer Queue
857 iter_except(s.pop, KeyError) # non-blocking set iterator
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000858
Raymond Hettinger30c73622010-12-01 23:45:20 +0000859 """
860 try:
861 if first is not None:
862 yield first() # For database APIs needing an initial cast to db.first()
Raymond Hettingera503f702016-03-13 00:12:31 -0800863 while True:
Raymond Hettinger30c73622010-12-01 23:45:20 +0000864 yield func()
865 except exception:
866 pass
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000867
Raymond Hettinger31b26f62014-04-02 03:16:42 -0700868 def first_true(iterable, default=False, pred=None):
869 """Returns the first true value in the iterable.
870
871 If no true value is found, returns *default*
872
873 If *pred* is not None, returns the first item
874 for which pred(item) is true.
875
876 """
877 # first_true([a,b,c], x) --> a or b or c or x
878 # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
879 return next(filter(pred, iterable), default)
880
Raymond Hettinger30c73622010-12-01 23:45:20 +0000881 def random_product(*args, repeat=1):
882 "Random selection from itertools.product(*args, **kwds)"
883 pools = [tuple(pool) for pool in args] * repeat
884 return tuple(random.choice(pool) for pool in pools)
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000885
Raymond Hettinger30c73622010-12-01 23:45:20 +0000886 def random_permutation(iterable, r=None):
887 "Random selection from itertools.permutations(iterable, r)"
888 pool = tuple(iterable)
889 r = len(pool) if r is None else r
890 return tuple(random.sample(pool, r))
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000891
Raymond Hettinger30c73622010-12-01 23:45:20 +0000892 def random_combination(iterable, r):
893 "Random selection from itertools.combinations(iterable, r)"
894 pool = tuple(iterable)
895 n = len(pool)
896 indices = sorted(random.sample(range(n), r))
897 return tuple(pool[i] for i in indices)
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000898
Raymond Hettinger30c73622010-12-01 23:45:20 +0000899 def random_combination_with_replacement(iterable, r):
900 "Random selection from itertools.combinations_with_replacement(iterable, r)"
901 pool = tuple(iterable)
902 n = len(pool)
903 indices = sorted(random.randrange(n) for i in range(r))
904 return tuple(pool[i] for i in indices)
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000905
Raymond Hettingerd37258d2018-01-13 10:35:40 -0800906 def nth_combination(iterable, r, index):
907 'Equivalent to list(combinations(iterable, r))[index]'
908 pool = tuple(iterable)
909 n = len(pool)
910 if r < 0 or r > n:
911 raise ValueError
912 c = 1
913 k = min(r, n-r)
914 for i in range(1, k+1):
915 c = c * (n - k + i) // i
916 if index < 0:
917 index += c
918 if index < 0 or index >= c:
919 raise IndexError
920 result = []
921 while r:
922 c, n, r = c*r//n, n-1, r-1
923 while index >= c:
924 index -= c
925 c, n = c*(n-r)//n, n-1
926 result.append(pool[-1-n])
927 return tuple(result)
928