blob: fd77f99a88f577e538594ad505d4d10e5554dad6 [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``
Raymond Hettingercc061d02020-11-30 20:42:54 -080058:func:`pairwise` iterable (p[0], p[1]), (p[1], p[2]) ``pairwise('ABCDEFG') --> AB BC CD DE EF FG``
Benjamin Peterson2989f582014-01-16 10:10:13 -050059:func:`starmap` func, seq func(\*seq[0]), func(\*seq[1]), ... ``starmap(pow, [(2,5), (3,2), (10,3)]) --> 32 9 1000``
60:func:`takewhile` pred, seq seq[0], seq[1], until pred fails ``takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4``
61:func:`tee` it, n it1, it2, ... itn splits one iterator into n
62:func:`zip_longest` p, q, ... (p[0], q[0]), (p[1], q[1]), ... ``zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-``
63============================ ============================ ================================================= =============================================================
Raymond Hettingerf76b9202009-02-17 20:00:59 +000064
Raymond Hettinger6693d7a2017-12-15 13:17:55 -080065**Combinatoric iterators:**
Raymond Hettingerf76b9202009-02-17 20:00:59 +000066
Raymond Hettinger7f587cd2009-04-10 19:43:50 +000067============================================== ==================== =============================================================
68Iterator Arguments Results
69============================================== ==================== =============================================================
70:func:`product` p, q, ... [repeat=1] cartesian product, equivalent to a nested for-loop
71:func:`permutations` p[, r] r-length tuples, all possible orderings, no repeated elements
Raymond Hettinger36c3c022009-11-19 01:20:07 +000072:func:`combinations` p, r r-length tuples, in sorted order, no repeated elements
73:func:`combinations_with_replacement` p, r r-length tuples, in sorted order, with repeated elements
Raymond Hettinger7f587cd2009-04-10 19:43:50 +000074============================================== ==================== =============================================================
Georg Brandl116aa622007-08-15 14:28:22 +000075
Benedikt Werner14e3c442019-03-21 16:28:49 +010076============================================== =============================================================
77Examples Results
78============================================== =============================================================
79``product('ABCD', repeat=2)`` ``AA AB AC AD BA BB BC BD CA CB CC CD DA DB DC DD``
80``permutations('ABCD', 2)`` ``AB AC AD BA BC BD CA CB CD DA DB DC``
81``combinations('ABCD', 2)`` ``AB AC AD BC BD CD``
82``combinations_with_replacement('ABCD'2)`` ``AA AB AC AD BB BC BD CC CD DD``
83============================================== =============================================================
84
Georg Brandl116aa622007-08-15 14:28:22 +000085
86.. _itertools-functions:
87
88Itertool functions
89------------------
90
91The following module functions all construct and return iterators. Some provide
92streams of infinite length, so they should only be accessed by functions or
93loops that truncate the stream.
94
Lisa Roach9718b592018-09-23 17:34:59 -070095.. function:: accumulate(iterable[, func, *, initial=None])
Raymond Hettingeradb81462010-12-01 22:50:36 +000096
Andrew Kuchling15b04eb2014-04-15 22:28:40 -040097 Make an iterator that returns accumulated sums, or accumulated
98 results of other binary functions (specified via the optional
Lisa Roach9718b592018-09-23 17:34:59 -070099 *func* argument).
100
101 If *func* is supplied, it should be a function
Andrew Kuchling15b04eb2014-04-15 22:28:40 -0400102 of two arguments. Elements of the input *iterable* may be any type
103 that can be accepted as arguments to *func*. (For example, with
104 the default operation of addition, elements may be any addable
105 type including :class:`~decimal.Decimal` or
Lisa Roach9718b592018-09-23 17:34:59 -0700106 :class:`~fractions.Fraction`.)
107
108 Usually, the number of elements output matches the input iterable.
109 However, if the keyword argument *initial* is provided, the
110 accumulation leads off with the *initial* value so that the output
111 has one more element than the input iterable.
Raymond Hettingeradb81462010-12-01 22:50:36 +0000112
Raymond Hettinger672866d2016-05-28 00:17:54 -0700113 Roughly equivalent to::
Raymond Hettinger5d446132011-03-27 18:52:10 -0700114
Lisa Roach9718b592018-09-23 17:34:59 -0700115 def accumulate(iterable, func=operator.add, *, initial=None):
Raymond Hettingeradb81462010-12-01 22:50:36 +0000116 'Return running totals'
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000117 # accumulate([1,2,3,4,5]) --> 1 3 6 10 15
Lisa Roach9718b592018-09-23 17:34:59 -0700118 # accumulate([1,2,3,4,5], initial=100) --> 100 101 103 106 110 115
Raymond Hettinger5d446132011-03-27 18:52:10 -0700119 # accumulate([1,2,3,4,5], operator.mul) --> 1 2 6 24 120
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000120 it = iter(iterable)
Lisa Roach9718b592018-09-23 17:34:59 -0700121 total = initial
122 if initial is None:
123 try:
124 total = next(it)
125 except StopIteration:
126 return
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000127 yield total
128 for element in it:
Raymond Hettinger5d446132011-03-27 18:52:10 -0700129 total = func(total, element)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000130 yield total
Raymond Hettingeradb81462010-12-01 22:50:36 +0000131
Raymond Hettinger295c1d42011-04-21 11:09:28 -0700132 There are a number of uses for the *func* argument. It can be set to
133 :func:`min` for a running minimum, :func:`max` for a running maximum, or
134 :func:`operator.mul` for a running product. Amortization tables can be
135 built by accumulating interest and applying payments. First-order
Georg Brandl5d941342016-02-26 19:37:12 +0100136 `recurrence relations <https://en.wikipedia.org/wiki/Recurrence_relation>`_
Raymond Hettinger295c1d42011-04-21 11:09:28 -0700137 can be modeled by supplying the initial value in the iterable and using only
138 the accumulated total in *func* argument::
Raymond Hettinger5d446132011-03-27 18:52:10 -0700139
140 >>> data = [3, 4, 6, 2, 1, 9, 0, 7, 5, 8]
141 >>> list(accumulate(data, operator.mul)) # running product
142 [3, 12, 72, 144, 144, 1296, 0, 0, 0, 0]
143 >>> list(accumulate(data, max)) # running maximum
144 [3, 4, 6, 6, 6, 9, 9, 9, 9, 9]
145
146 # Amortize a 5% loan of 1000 with 4 annual payments of 90
147 >>> cashflows = [1000, -90, -90, -90, -90]
148 >>> list(accumulate(cashflows, lambda bal, pmt: bal*1.05 + pmt))
149 [1000, 960.0, 918.0, 873.9000000000001, 827.5950000000001]
150
Georg Brandl5d941342016-02-26 19:37:12 +0100151 # Chaotic recurrence relation https://en.wikipedia.org/wiki/Logistic_map
Raymond Hettinger295c1d42011-04-21 11:09:28 -0700152 >>> logistic_map = lambda x, _: r * x * (1 - x)
153 >>> r = 3.8
154 >>> x0 = 0.4
155 >>> inputs = repeat(x0, 36) # only the initial value is used
156 >>> [format(x, '.2f') for x in accumulate(inputs, logistic_map)]
157 ['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 +0200158 '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 -0700159 '0.93', '0.25', '0.71', '0.79', '0.63', '0.88', '0.39', '0.91', '0.32',
160 '0.83', '0.54', '0.95', '0.20', '0.60', '0.91', '0.30', '0.80', '0.60']
161
Raymond Hettinger64801682013-10-12 16:04:17 -0700162 See :func:`functools.reduce` for a similar function that returns only the
163 final accumulated value.
164
Raymond Hettingeradb81462010-12-01 22:50:36 +0000165 .. versionadded:: 3.2
Georg Brandl116aa622007-08-15 14:28:22 +0000166
Raymond Hettinger5d446132011-03-27 18:52:10 -0700167 .. versionchanged:: 3.3
168 Added the optional *func* parameter.
169
Lisa Roach9718b592018-09-23 17:34:59 -0700170 .. versionchanged:: 3.8
171 Added the optional *initial* parameter.
172
Georg Brandl116aa622007-08-15 14:28:22 +0000173.. function:: chain(*iterables)
174
Raymond Hettinger30c73622010-12-01 23:45:20 +0000175 Make an iterator that returns elements from the first iterable until it is
176 exhausted, then proceeds to the next iterable, until all of the iterables are
177 exhausted. Used for treating consecutive sequences as a single sequence.
Raymond Hettinger672866d2016-05-28 00:17:54 -0700178 Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000179
Raymond Hettinger30c73622010-12-01 23:45:20 +0000180 def chain(*iterables):
181 # chain('ABC', 'DEF') --> A B C D E F
182 for it in iterables:
183 for element in it:
184 yield element
Georg Brandl116aa622007-08-15 14:28:22 +0000185
186
Georg Brandl933b9742010-07-29 14:36:11 +0000187.. classmethod:: chain.from_iterable(iterable)
Christian Heimes70e7ea22008-02-28 20:02:27 +0000188
Raymond Hettinger30c73622010-12-01 23:45:20 +0000189 Alternate constructor for :func:`chain`. Gets chained inputs from a
Raymond Hettinger1e21ebc2013-09-09 01:54:27 -0500190 single iterable argument that is evaluated lazily. Roughly equivalent to::
Christian Heimes70e7ea22008-02-28 20:02:27 +0000191
Raymond Hettinger30c73622010-12-01 23:45:20 +0000192 def from_iterable(iterables):
193 # chain.from_iterable(['ABC', 'DEF']) --> A B C D E F
194 for it in iterables:
195 for element in it:
196 yield element
Christian Heimes70e7ea22008-02-28 20:02:27 +0000197
Christian Heimes78644762008-03-04 23:39:23 +0000198
Christian Heimes836baa52008-02-26 08:18:30 +0000199.. function:: combinations(iterable, r)
200
Raymond Hettinger30c73622010-12-01 23:45:20 +0000201 Return *r* length subsequences of elements from the input *iterable*.
Christian Heimes836baa52008-02-26 08:18:30 +0000202
Ruaridh Williamson5e0ed8a2020-05-28 20:56:43 +0100203 The combination tuples are emitted in lexicographic ordering according to
204 the order of the input *iterable*. So, if the input *iterable* is sorted,
205 the combination tuples will be produced in sorted order.
Christian Heimes836baa52008-02-26 08:18:30 +0000206
Raymond Hettinger30c73622010-12-01 23:45:20 +0000207 Elements are treated as unique based on their position, not on their
208 value. So if the input elements are unique, there will be no repeat
209 values in each combination.
Christian Heimes836baa52008-02-26 08:18:30 +0000210
Raymond Hettinger672866d2016-05-28 00:17:54 -0700211 Roughly equivalent to::
Christian Heimes836baa52008-02-26 08:18:30 +0000212
213 def combinations(iterable, r):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000214 # combinations('ABCD', 2) --> AB AC AD BC BD CD
215 # combinations(range(4), 3) --> 012 013 023 123
Christian Heimes836baa52008-02-26 08:18:30 +0000216 pool = tuple(iterable)
Christian Heimes380f7f22008-02-28 11:19:05 +0000217 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000218 if r > n:
219 return
220 indices = list(range(r))
Christian Heimesb558a2e2008-03-02 22:46:37 +0000221 yield tuple(pool[i] for i in indices)
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000222 while True:
Christian Heimes380f7f22008-02-28 11:19:05 +0000223 for i in reversed(range(r)):
Christian Heimesb558a2e2008-03-02 22:46:37 +0000224 if indices[i] != i + n - r:
Christian Heimes836baa52008-02-26 08:18:30 +0000225 break
Christian Heimes380f7f22008-02-28 11:19:05 +0000226 else:
227 return
Christian Heimesb558a2e2008-03-02 22:46:37 +0000228 indices[i] += 1
Christian Heimes380f7f22008-02-28 11:19:05 +0000229 for j in range(i+1, r):
Christian Heimesb558a2e2008-03-02 22:46:37 +0000230 indices[j] = indices[j-1] + 1
231 yield tuple(pool[i] for i in indices)
Christian Heimes836baa52008-02-26 08:18:30 +0000232
Raymond Hettinger30c73622010-12-01 23:45:20 +0000233 The code for :func:`combinations` can be also expressed as a subsequence
234 of :func:`permutations` after filtering entries where the elements are not
235 in sorted order (according to their position in the input pool)::
Christian Heimes78644762008-03-04 23:39:23 +0000236
237 def combinations(iterable, r):
238 pool = tuple(iterable)
239 n = len(pool)
240 for indices in permutations(range(n), r):
241 if sorted(indices) == list(indices):
242 yield tuple(pool[i] for i in indices)
243
Raymond Hettinger30c73622010-12-01 23:45:20 +0000244 The number of items returned is ``n! / r! / (n-r)!`` when ``0 <= r <= n``
245 or zero when ``r > n``.
Christian Heimes836baa52008-02-26 08:18:30 +0000246
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000247.. function:: combinations_with_replacement(iterable, r)
248
Raymond Hettinger30c73622010-12-01 23:45:20 +0000249 Return *r* length subsequences of elements from the input *iterable*
250 allowing individual elements to be repeated more than once.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000251
Ruaridh Williamson5e0ed8a2020-05-28 20:56:43 +0100252 The combination tuples are emitted in lexicographic ordering according to
253 the order of the input *iterable*. So, if the input *iterable* is sorted,
254 the combination tuples will be produced in sorted order.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000255
Raymond Hettinger30c73622010-12-01 23:45:20 +0000256 Elements are treated as unique based on their position, not on their
257 value. So if the input elements are unique, the generated combinations
258 will also be unique.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000259
Raymond Hettinger672866d2016-05-28 00:17:54 -0700260 Roughly equivalent to::
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000261
262 def combinations_with_replacement(iterable, r):
263 # combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC
264 pool = tuple(iterable)
265 n = len(pool)
266 if not n and r:
267 return
268 indices = [0] * r
269 yield tuple(pool[i] for i in indices)
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000270 while True:
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000271 for i in reversed(range(r)):
272 if indices[i] != n - 1:
273 break
274 else:
275 return
276 indices[i:] = [indices[i] + 1] * (r - i)
277 yield tuple(pool[i] for i in indices)
278
Raymond Hettinger30c73622010-12-01 23:45:20 +0000279 The code for :func:`combinations_with_replacement` can be also expressed as
280 a subsequence of :func:`product` after filtering entries where the elements
281 are not in sorted order (according to their position in the input pool)::
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000282
283 def combinations_with_replacement(iterable, r):
284 pool = tuple(iterable)
285 n = len(pool)
286 for indices in product(range(n), repeat=r):
287 if sorted(indices) == list(indices):
288 yield tuple(pool[i] for i in indices)
289
Raymond Hettinger30c73622010-12-01 23:45:20 +0000290 The number of items returned is ``(n+r-1)! / r! / (n-1)!`` when ``n > 0``.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000291
Raymond Hettinger30c73622010-12-01 23:45:20 +0000292 .. versionadded:: 3.1
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000293
Georg Brandl67b21b72010-08-17 15:07:14 +0000294
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000295.. function:: compress(data, selectors)
296
Raymond Hettinger30c73622010-12-01 23:45:20 +0000297 Make an iterator that filters elements from *data* returning only those that
298 have a corresponding element in *selectors* that evaluates to ``True``.
299 Stops when either the *data* or *selectors* iterables has been exhausted.
Raymond Hettinger672866d2016-05-28 00:17:54 -0700300 Roughly equivalent to::
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000301
Raymond Hettinger30c73622010-12-01 23:45:20 +0000302 def compress(data, selectors):
303 # compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F
304 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000305
Raymond Hettinger30c73622010-12-01 23:45:20 +0000306 .. versionadded:: 3.1
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000307
308
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000309.. function:: count(start=0, step=1)
Georg Brandl116aa622007-08-15 14:28:22 +0000310
Andrew Kuchlingedb42602013-06-21 08:00:58 -0400311 Make an iterator that returns evenly spaced values starting with number *start*. Often
Raymond Hettinger30c73622010-12-01 23:45:20 +0000312 used as an argument to :func:`map` to generate consecutive data points.
Raymond Hettinger672866d2016-05-28 00:17:54 -0700313 Also, used with :func:`zip` to add sequence numbers. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000314
Raymond Hettinger30c73622010-12-01 23:45:20 +0000315 def count(start=0, step=1):
316 # count(10) --> 10 11 12 13 14 ...
Georg Brandl37a80dc2011-01-13 07:31:18 +0000317 # count(2.5, 0.5) -> 2.5 3.0 3.5 ...
Raymond Hettinger30c73622010-12-01 23:45:20 +0000318 n = start
319 while True:
320 yield n
321 n += step
Georg Brandl116aa622007-08-15 14:28:22 +0000322
Raymond Hettinger30c73622010-12-01 23:45:20 +0000323 When counting with floating point numbers, better accuracy can sometimes be
324 achieved by substituting multiplicative code such as: ``(start + step * i
325 for i in count())``.
Raymond Hettinger5bc472a2009-06-17 01:40:52 +0000326
Raymond Hettinger30c73622010-12-01 23:45:20 +0000327 .. versionchanged:: 3.1
328 Added *step* argument and allowed non-integer arguments.
Georg Brandl116aa622007-08-15 14:28:22 +0000329
330.. function:: cycle(iterable)
331
Raymond Hettinger30c73622010-12-01 23:45:20 +0000332 Make an iterator returning elements from the iterable and saving a copy of each.
333 When the iterable is exhausted, return elements from the saved copy. Repeats
Raymond Hettinger672866d2016-05-28 00:17:54 -0700334 indefinitely. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000335
Raymond Hettinger30c73622010-12-01 23:45:20 +0000336 def cycle(iterable):
337 # cycle('ABCD') --> A B C D A B C D A B C D ...
338 saved = []
339 for element in iterable:
340 yield element
341 saved.append(element)
342 while saved:
343 for element in saved:
Georg Brandl116aa622007-08-15 14:28:22 +0000344 yield element
345
Raymond Hettinger30c73622010-12-01 23:45:20 +0000346 Note, this member of the toolkit may require significant auxiliary storage
347 (depending on the length of the iterable).
Georg Brandl116aa622007-08-15 14:28:22 +0000348
349
350.. function:: dropwhile(predicate, iterable)
351
Raymond Hettinger30c73622010-12-01 23:45:20 +0000352 Make an iterator that drops elements from the iterable as long as the predicate
353 is true; afterwards, returns every element. Note, the iterator does not produce
354 *any* output until the predicate first becomes false, so it may have a lengthy
Raymond Hettinger672866d2016-05-28 00:17:54 -0700355 start-up time. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000356
Raymond Hettinger30c73622010-12-01 23:45:20 +0000357 def dropwhile(predicate, iterable):
358 # dropwhile(lambda x: x<5, [1,4,6,4,1]) --> 6 4 1
359 iterable = iter(iterable)
360 for x in iterable:
361 if not predicate(x):
362 yield x
363 break
364 for x in iterable:
365 yield x
Georg Brandl116aa622007-08-15 14:28:22 +0000366
Raymond Hettinger749761e2009-01-27 04:42:48 +0000367.. function:: filterfalse(predicate, iterable)
368
Raymond Hettinger30c73622010-12-01 23:45:20 +0000369 Make an iterator that filters elements from iterable returning only those for
370 which the predicate is ``False``. If *predicate* is ``None``, return the items
Raymond Hettinger672866d2016-05-28 00:17:54 -0700371 that are false. Roughly equivalent to::
Raymond Hettinger749761e2009-01-27 04:42:48 +0000372
Raymond Hettinger30c73622010-12-01 23:45:20 +0000373 def filterfalse(predicate, iterable):
374 # filterfalse(lambda x: x%2, range(10)) --> 0 2 4 6 8
375 if predicate is None:
376 predicate = bool
377 for x in iterable:
378 if not predicate(x):
379 yield x
Raymond Hettinger749761e2009-01-27 04:42:48 +0000380
Georg Brandl116aa622007-08-15 14:28:22 +0000381
Georg Brandl3dd33882009-06-01 17:35:27 +0000382.. function:: groupby(iterable, key=None)
Georg Brandl116aa622007-08-15 14:28:22 +0000383
Raymond Hettinger30c73622010-12-01 23:45:20 +0000384 Make an iterator that returns consecutive keys and groups from the *iterable*.
385 The *key* is a function computing a key value for each element. If not
386 specified or is ``None``, *key* defaults to an identity function and returns
387 the element unchanged. Generally, the iterable needs to already be sorted on
388 the same key function.
Georg Brandl116aa622007-08-15 14:28:22 +0000389
Raymond Hettinger30c73622010-12-01 23:45:20 +0000390 The operation of :func:`groupby` is similar to the ``uniq`` filter in Unix. It
391 generates a break or new group every time the value of the key function changes
392 (which is why it is usually necessary to have sorted the data using the same key
393 function). That behavior differs from SQL's GROUP BY which aggregates common
394 elements regardless of their input order.
Georg Brandl116aa622007-08-15 14:28:22 +0000395
Raymond Hettinger30c73622010-12-01 23:45:20 +0000396 The returned group is itself an iterator that shares the underlying iterable
397 with :func:`groupby`. Because the source is shared, when the :func:`groupby`
398 object is advanced, the previous group is no longer visible. So, if that data
399 is needed later, it should be stored as a list::
Georg Brandl116aa622007-08-15 14:28:22 +0000400
Raymond Hettinger30c73622010-12-01 23:45:20 +0000401 groups = []
402 uniquekeys = []
403 data = sorted(data, key=keyfunc)
404 for k, g in groupby(data, keyfunc):
405 groups.append(list(g)) # Store group iterator as a list
406 uniquekeys.append(k)
Georg Brandl116aa622007-08-15 14:28:22 +0000407
Raymond Hettinger672866d2016-05-28 00:17:54 -0700408 :func:`groupby` is roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000409
Raymond Hettinger30c73622010-12-01 23:45:20 +0000410 class groupby:
411 # [k for k, g in groupby('AAAABBBCCDAABBB')] --> A B C D A B
412 # [list(g) for k, g in groupby('AAAABBBCCD')] --> AAAA BBB CC D
413 def __init__(self, iterable, key=None):
414 if key is None:
415 key = lambda x: x
416 self.keyfunc = key
417 self.it = iter(iterable)
418 self.tgtkey = self.currkey = self.currvalue = object()
419 def __iter__(self):
420 return self
421 def __next__(self):
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300422 self.id = object()
Raymond Hettinger30c73622010-12-01 23:45:20 +0000423 while self.currkey == self.tgtkey:
424 self.currvalue = next(self.it) # Exit on StopIteration
425 self.currkey = self.keyfunc(self.currvalue)
426 self.tgtkey = self.currkey
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300427 return (self.currkey, self._grouper(self.tgtkey, self.id))
428 def _grouper(self, tgtkey, id):
429 while self.id is id and self.currkey == tgtkey:
Raymond Hettinger30c73622010-12-01 23:45:20 +0000430 yield self.currvalue
Raymond Hettinger828d9322014-11-22 21:56:23 -0800431 try:
432 self.currvalue = next(self.it)
433 except StopIteration:
434 return
Raymond Hettinger30c73622010-12-01 23:45:20 +0000435 self.currkey = self.keyfunc(self.currvalue)
Georg Brandl116aa622007-08-15 14:28:22 +0000436
Georg Brandl116aa622007-08-15 14:28:22 +0000437
Ezio Melottie0add762012-09-14 06:32:35 +0300438.. function:: islice(iterable, stop)
439 islice(iterable, start, stop[, step])
Georg Brandl116aa622007-08-15 14:28:22 +0000440
Raymond Hettinger30c73622010-12-01 23:45:20 +0000441 Make an iterator that returns selected elements from the iterable. If *start* is
442 non-zero, then elements from the iterable are skipped until start is reached.
443 Afterward, elements are returned consecutively unless *step* is set higher than
444 one which results in items being skipped. If *stop* is ``None``, then iteration
445 continues until the iterator is exhausted, if at all; otherwise, it stops at the
446 specified position. Unlike regular slicing, :func:`islice` does not support
447 negative values for *start*, *stop*, or *step*. Can be used to extract related
448 fields from data where the internal structure has been flattened (for example, a
Raymond Hettinger672866d2016-05-28 00:17:54 -0700449 multi-line report may list a name field on every third line). Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000450
Raymond Hettinger30c73622010-12-01 23:45:20 +0000451 def islice(iterable, *args):
452 # islice('ABCDEFG', 2) --> A B
453 # islice('ABCDEFG', 2, 4) --> C D
454 # islice('ABCDEFG', 2, None) --> C D E F G
455 # islice('ABCDEFG', 0, None, 2) --> A C E G
456 s = slice(*args)
Cheryl Sabellada1734c2018-03-26 21:29:33 -0400457 start, stop, step = s.start or 0, s.stop or sys.maxsize, s.step or 1
458 it = iter(range(start, stop, step))
Raymond Hettinger828d9322014-11-22 21:56:23 -0800459 try:
460 nexti = next(it)
461 except StopIteration:
Cheryl Sabellada1734c2018-03-26 21:29:33 -0400462 # Consume *iterable* up to the *start* position.
463 for i, element in zip(range(start), iterable):
464 pass
Raymond Hettinger828d9322014-11-22 21:56:23 -0800465 return
Cheryl Sabellada1734c2018-03-26 21:29:33 -0400466 try:
467 for i, element in enumerate(iterable):
468 if i == nexti:
469 yield element
470 nexti = next(it)
471 except StopIteration:
472 # Consume to *stop*.
473 for i, element in zip(range(i + 1, stop), iterable):
474 pass
Georg Brandl116aa622007-08-15 14:28:22 +0000475
Raymond Hettinger30c73622010-12-01 23:45:20 +0000476 If *start* is ``None``, then iteration starts at zero. If *step* is ``None``,
477 then the step defaults to one.
Georg Brandl116aa622007-08-15 14:28:22 +0000478
Raymond Hettingercc061d02020-11-30 20:42:54 -0800479.. function:: pairwise(iterable)
480
481 Return successive overlapping pairs taken from the input *iterable*.
482
483 The number of 2-tuples in the output iterator will be one fewer than the
484 number of inputs. It will be empty if the input iterable has fewer than
485 two values.
486
487 Roughly equivalent to::
488
489 def pairwise(iterable):
490 # pairwise('ABCDEFG') --> AB BC CD DE EF FG
491 a, b = tee(iterable)
492 next(b, None)
493 return zip(a, b)
494
Georg Brandl116aa622007-08-15 14:28:22 +0000495
Georg Brandl3dd33882009-06-01 17:35:27 +0000496.. function:: permutations(iterable, r=None)
Christian Heimes70e7ea22008-02-28 20:02:27 +0000497
Raymond Hettinger30c73622010-12-01 23:45:20 +0000498 Return successive *r* length permutations of elements in the *iterable*.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000499
Raymond Hettinger30c73622010-12-01 23:45:20 +0000500 If *r* is not specified or is ``None``, then *r* defaults to the length
501 of the *iterable* and all possible full-length permutations
502 are generated.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000503
Ruaridh Williamson5e0ed8a2020-05-28 20:56:43 +0100504 The permutation tuples are emitted in lexicographic ordering according to
505 the order of the input *iterable*. So, if the input *iterable* is sorted,
506 the combination tuples will be produced in sorted order.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000507
Raymond Hettinger30c73622010-12-01 23:45:20 +0000508 Elements are treated as unique based on their position, not on their
509 value. So if the input elements are unique, there will be no repeat
510 values in each permutation.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000511
Raymond Hettinger672866d2016-05-28 00:17:54 -0700512 Roughly equivalent to::
Christian Heimesb558a2e2008-03-02 22:46:37 +0000513
514 def permutations(iterable, r=None):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000515 # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
516 # permutations(range(3)) --> 012 021 102 120 201 210
Christian Heimesb558a2e2008-03-02 22:46:37 +0000517 pool = tuple(iterable)
518 n = len(pool)
519 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000520 if r > n:
521 return
522 indices = list(range(n))
Sandro Tosi73866622011-12-25 17:25:45 +0100523 cycles = list(range(n, n-r, -1))
Christian Heimesb558a2e2008-03-02 22:46:37 +0000524 yield tuple(pool[i] for i in indices[:r])
525 while n:
526 for i in reversed(range(r)):
527 cycles[i] -= 1
528 if cycles[i] == 0:
529 indices[i:] = indices[i+1:] + indices[i:i+1]
530 cycles[i] = n - i
531 else:
532 j = cycles[i]
533 indices[i], indices[-j] = indices[-j], indices[i]
534 yield tuple(pool[i] for i in indices[:r])
535 break
536 else:
537 return
Christian Heimes70e7ea22008-02-28 20:02:27 +0000538
Raymond Hettinger30c73622010-12-01 23:45:20 +0000539 The code for :func:`permutations` can be also expressed as a subsequence of
540 :func:`product`, filtered to exclude entries with repeated elements (those
541 from the same position in the input pool)::
Christian Heimes78644762008-03-04 23:39:23 +0000542
543 def permutations(iterable, r=None):
544 pool = tuple(iterable)
545 n = len(pool)
546 r = n if r is None else r
547 for indices in product(range(n), repeat=r):
548 if len(set(indices)) == r:
549 yield tuple(pool[i] for i in indices)
550
Raymond Hettinger30c73622010-12-01 23:45:20 +0000551 The number of items returned is ``n! / (n-r)!`` when ``0 <= r <= n``
552 or zero when ``r > n``.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000553
Georg Brandl3dd33882009-06-01 17:35:27 +0000554.. function:: product(*iterables, repeat=1)
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000555
Raymond Hettinger30c73622010-12-01 23:45:20 +0000556 Cartesian product of input iterables.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000557
Raymond Hettinger672866d2016-05-28 00:17:54 -0700558 Roughly equivalent to nested for-loops in a generator expression. For example,
Raymond Hettinger30c73622010-12-01 23:45:20 +0000559 ``product(A, B)`` returns the same as ``((x,y) for x in A for y in B)``.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000560
Raymond Hettinger30c73622010-12-01 23:45:20 +0000561 The nested loops cycle like an odometer with the rightmost element advancing
562 on every iteration. This pattern creates a lexicographic ordering so that if
563 the input's iterables are sorted, the product tuples are emitted in sorted
564 order.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000565
Raymond Hettinger30c73622010-12-01 23:45:20 +0000566 To compute the product of an iterable with itself, specify the number of
567 repetitions with the optional *repeat* keyword argument. For example,
568 ``product(A, repeat=4)`` means the same as ``product(A, A, A, A)``.
Christian Heimes9e7f1d22008-02-28 12:27:11 +0000569
Raymond Hettinger672866d2016-05-28 00:17:54 -0700570 This function is roughly equivalent to the following code, except that the
Raymond Hettinger30c73622010-12-01 23:45:20 +0000571 actual implementation does not build up intermediate results in memory::
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000572
Raymond Hettinger30c73622010-12-01 23:45:20 +0000573 def product(*args, repeat=1):
574 # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy
575 # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111
576 pools = [tuple(pool) for pool in args] * repeat
577 result = [[]]
578 for pool in pools:
579 result = [x+[y] for x in result for y in pool]
580 for prod in result:
581 yield tuple(prod)
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000582
Ramil Nugmanovcfc6ce42020-05-28 19:46:22 +0300583 Before :func:`product` runs, it completely consumes the input iterables,
584 keeping pools of values in memory to generate the products. Accordingly,
Zackery Spytz074ad512020-12-17 13:24:43 -0700585 it is only useful with finite inputs.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000586
Raymond Hettingerd75ad442009-06-01 19:16:52 +0000587.. function:: repeat(object[, times])
Georg Brandl116aa622007-08-15 14:28:22 +0000588
Raymond Hettinger30c73622010-12-01 23:45:20 +0000589 Make an iterator that returns *object* over and over again. Runs indefinitely
590 unless the *times* argument is specified. Used as argument to :func:`map` for
591 invariant parameters to the called function. Also used with :func:`zip` to
Raymond Hettinger672866d2016-05-28 00:17:54 -0700592 create an invariant part of a tuple record.
593
594 Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000595
Raymond Hettinger30c73622010-12-01 23:45:20 +0000596 def repeat(object, times=None):
597 # repeat(10, 3) --> 10 10 10
598 if times is None:
599 while True:
600 yield object
601 else:
602 for i in range(times):
603 yield object
Georg Brandl116aa622007-08-15 14:28:22 +0000604
Raymond Hettingerfc3ba6b2012-02-01 09:07:40 -0800605 A common use for *repeat* is to supply a stream of constant values to *map*
606 or *zip*::
607
608 >>> list(map(pow, range(10), repeat(2)))
609 [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
Georg Brandl116aa622007-08-15 14:28:22 +0000610
611.. function:: starmap(function, iterable)
612
Raymond Hettinger30c73622010-12-01 23:45:20 +0000613 Make an iterator that computes the function using arguments obtained from
614 the iterable. Used instead of :func:`map` when argument parameters are already
615 grouped in tuples from a single iterable (the data has been "pre-zipped"). The
616 difference between :func:`map` and :func:`starmap` parallels the distinction
Raymond Hettinger672866d2016-05-28 00:17:54 -0700617 between ``function(a,b)`` and ``function(*c)``. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000618
Raymond Hettinger30c73622010-12-01 23:45:20 +0000619 def starmap(function, iterable):
620 # starmap(pow, [(2,5), (3,2), (10,3)]) --> 32 9 1000
621 for args in iterable:
622 yield function(*args)
Christian Heimes679db4a2008-01-18 09:56:22 +0000623
Georg Brandl116aa622007-08-15 14:28:22 +0000624
625.. function:: takewhile(predicate, iterable)
626
Raymond Hettinger30c73622010-12-01 23:45:20 +0000627 Make an iterator that returns elements from the iterable as long as the
Raymond Hettinger672866d2016-05-28 00:17:54 -0700628 predicate is true. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000629
Raymond Hettinger30c73622010-12-01 23:45:20 +0000630 def takewhile(predicate, iterable):
631 # takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4
632 for x in iterable:
633 if predicate(x):
634 yield x
635 else:
636 break
Georg Brandl116aa622007-08-15 14:28:22 +0000637
638
Georg Brandl3dd33882009-06-01 17:35:27 +0000639.. function:: tee(iterable, n=2)
Georg Brandl116aa622007-08-15 14:28:22 +0000640
Raymond Hettingerab425aa2016-04-26 00:10:00 -0700641 Return *n* independent iterators from a single iterable.
642
643 The following Python code helps explain what *tee* does (although the actual
Serhiy Storchaka4ecfa452016-05-16 09:31:54 +0300644 implementation is more complex and uses only a single underlying
Raymond Hettinger672866d2016-05-28 00:17:54 -0700645 :abbr:`FIFO (first-in, first-out)` queue).
646
647 Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000648
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000649 def tee(iterable, n=2):
650 it = iter(iterable)
651 deques = [collections.deque() for i in range(n)]
652 def gen(mydeque):
653 while True:
654 if not mydeque: # when the local deque is empty
Raymond Hettinger828d9322014-11-22 21:56:23 -0800655 try:
656 newval = next(it) # fetch a new value and
657 except StopIteration:
658 return
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000659 for d in deques: # load it to all the deques
660 d.append(newval)
661 yield mydeque.popleft()
662 return tuple(gen(d) for d in deques)
Georg Brandl116aa622007-08-15 14:28:22 +0000663
Raymond Hettinger30c73622010-12-01 23:45:20 +0000664 Once :func:`tee` has made a split, the original *iterable* should not be
665 used anywhere else; otherwise, the *iterable* could get advanced without
Serhiy Storchaka918b4682019-09-09 11:18:16 +0300666 the tee objects being informed.
Georg Brandl116aa622007-08-15 14:28:22 +0000667
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300668 ``tee`` iterators are not threadsafe. A :exc:`RuntimeError` may be
669 raised when using simultaneously iterators returned by the same :func:`tee`
670 call, even if the original *iterable* is threadsafe.
671
Raymond Hettinger30c73622010-12-01 23:45:20 +0000672 This itertool may require significant auxiliary storage (depending on how
673 much temporary data needs to be stored). In general, if one iterator uses
674 most or all of the data before another iterator starts, it is faster to use
675 :func:`list` instead of :func:`tee`.
Georg Brandl116aa622007-08-15 14:28:22 +0000676
Georg Brandl116aa622007-08-15 14:28:22 +0000677
Georg Brandl3dd33882009-06-01 17:35:27 +0000678.. function:: zip_longest(*iterables, fillvalue=None)
Raymond Hettinger749761e2009-01-27 04:42:48 +0000679
Raymond Hettinger30c73622010-12-01 23:45:20 +0000680 Make an iterator that aggregates elements from each of the iterables. If the
681 iterables are of uneven length, missing values are filled-in with *fillvalue*.
Raymond Hettinger672866d2016-05-28 00:17:54 -0700682 Iteration continues until the longest iterable is exhausted. Roughly equivalent to::
Raymond Hettinger749761e2009-01-27 04:42:48 +0000683
Raymond Hettinger3147b042017-09-07 14:01:44 -0700684 def zip_longest(*args, fillvalue=None):
Raymond Hettinger30c73622010-12-01 23:45:20 +0000685 # zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-
Raymond Hettinger3147b042017-09-07 14:01:44 -0700686 iterators = [iter(it) for it in args]
687 num_active = len(iterators)
688 if not num_active:
689 return
690 while True:
691 values = []
692 for i, it in enumerate(iterators):
693 try:
694 value = next(it)
695 except StopIteration:
696 num_active -= 1
697 if not num_active:
698 return
699 iterators[i] = repeat(fillvalue)
700 value = fillvalue
701 values.append(value)
702 yield tuple(values)
Raymond Hettinger749761e2009-01-27 04:42:48 +0000703
Raymond Hettinger30c73622010-12-01 23:45:20 +0000704 If one of the iterables is potentially infinite, then the :func:`zip_longest`
705 function should be wrapped with something that limits the number of calls
706 (for example :func:`islice` or :func:`takewhile`). If not specified,
707 *fillvalue* defaults to ``None``.
Raymond Hettinger749761e2009-01-27 04:42:48 +0000708
709
Georg Brandl116aa622007-08-15 14:28:22 +0000710.. _itertools-recipes:
711
Raymond Hettinger1fa76822010-12-06 23:31:36 +0000712Itertools Recipes
713-----------------
Georg Brandl116aa622007-08-15 14:28:22 +0000714
715This section shows recipes for creating an extended toolset using the existing
716itertools as building blocks.
717
Raymond Hettingeradf02b32019-08-04 13:35:58 -0700718Substantially all of these recipes and many, many others can be installed from
719the `more-itertools project <https://pypi.org/project/more-itertools/>`_ found
720on the Python Package Index::
721
722 pip install more-itertools
723
Georg Brandl116aa622007-08-15 14:28:22 +0000724The extended tools offer the same high performance as the underlying toolset.
725The superior memory performance is kept by processing elements one at a time
726rather than bringing the whole iterable into memory all at once. Code volume is
727kept small by linking the tools together in a functional style which helps
728eliminate temporary variables. High speed is retained by preferring
Georg Brandl9afde1c2007-11-01 20:32:30 +0000729"vectorized" building blocks over the use of for-loops and :term:`generator`\s
Christian Heimesfe337bf2008-03-23 21:54:12 +0000730which incur interpreter overhead.
731
732.. testcode::
Georg Brandl116aa622007-08-15 14:28:22 +0000733
Raymond Hettinger30c73622010-12-01 23:45:20 +0000734 def take(n, iterable):
735 "Return first n items of the iterable as a list"
736 return list(islice(iterable, n))
Georg Brandl116aa622007-08-15 14:28:22 +0000737
Raymond Hettinger9265dd72018-04-08 08:44:20 -0700738 def prepend(value, iterator):
739 "Prepend a single value in front of an iterator"
740 # prepend(1, [2, 3, 4]) -> 1 2 3 4
741 return chain([value], iterator)
742
Raymond Hettinger30c73622010-12-01 23:45:20 +0000743 def tabulate(function, start=0):
744 "Return function(0), function(1), ..."
745 return map(function, count(start))
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000746
Raymond Hettingerdfe098d2014-05-25 22:03:56 -0700747 def tail(n, iterable):
748 "Return an iterator over the last n items"
749 # tail(3, 'ABCDEFG') --> E F G
750 return iter(collections.deque(iterable, maxlen=n))
751
Cheryl Sabellada1734c2018-03-26 21:29:33 -0400752 def consume(iterator, n=None):
753 "Advance the iterator n-steps ahead. If n is None, consume entirely."
Raymond Hettinger30c73622010-12-01 23:45:20 +0000754 # Use functions that consume iterators at C speed.
755 if n is None:
756 # feed the entire iterator into a zero-length deque
757 collections.deque(iterator, maxlen=0)
758 else:
759 # advance to the empty slice starting at position n
760 next(islice(iterator, n, n), None)
Raymond Hettingerfa007962009-03-09 11:55:25 +0000761
Raymond Hettinger30c73622010-12-01 23:45:20 +0000762 def nth(iterable, n, default=None):
763 "Returns the nth item or a default value"
764 return next(islice(iterable, n, None), default)
Georg Brandl116aa622007-08-15 14:28:22 +0000765
Raymond Hettingere525ee32016-03-06 18:11:38 -0800766 def all_equal(iterable):
767 "Returns True if all the elements are equal to each other"
768 g = groupby(iterable)
769 return next(g, True) and not next(g, False)
770
Raymond Hettinger30c73622010-12-01 23:45:20 +0000771 def quantify(iterable, pred=bool):
772 "Count how many times the predicate is true"
773 return sum(map(pred, iterable))
Georg Brandl116aa622007-08-15 14:28:22 +0000774
Raymond Hettingerfc40b302020-11-29 10:47:22 -0800775 def pad_none(iterable):
Raymond Hettinger30c73622010-12-01 23:45:20 +0000776 """Returns the sequence elements and then returns None indefinitely.
Georg Brandl116aa622007-08-15 14:28:22 +0000777
Raymond Hettinger30c73622010-12-01 23:45:20 +0000778 Useful for emulating the behavior of the built-in map() function.
779 """
780 return chain(iterable, repeat(None))
Georg Brandl116aa622007-08-15 14:28:22 +0000781
Raymond Hettinger30c73622010-12-01 23:45:20 +0000782 def ncycles(iterable, n):
783 "Returns the sequence elements n times"
784 return chain.from_iterable(repeat(tuple(iterable), n))
Georg Brandl116aa622007-08-15 14:28:22 +0000785
Raymond Hettinger30c73622010-12-01 23:45:20 +0000786 def dotproduct(vec1, vec2):
787 return sum(map(operator.mul, vec1, vec2))
Georg Brandl116aa622007-08-15 14:28:22 +0000788
Raymond Hettinger77fde8d2020-12-25 16:43:20 -0800789 def convolve(signal, kernel):
790 # See: https://betterexplained.com/articles/intuitive-convolution/
791 # convolve(data, [0.25, 0.25, 0.25, 0.25]) --> Moving average (blur)
792 # convolve(data, [1, -1]) --> 1st finite difference (1st derivative)
793 # convolve(data, [1, -2, 1]) --> 2nd finite difference (2nd derivative)
Raymond Hettingerf421bfc2020-12-30 12:51:19 -0800794 kernel = tuple(kernel)[::-1]
Raymond Hettinger77fde8d2020-12-25 16:43:20 -0800795 n = len(kernel)
Raymond Hettingerf421bfc2020-12-30 12:51:19 -0800796 window = collections.deque([0], maxlen=n) * n
Raymond Hettinger77fde8d2020-12-25 16:43:20 -0800797 for x in chain(signal, repeat(0, n-1)):
798 window.append(x)
799 yield sum(map(operator.mul, kernel, window))
800
Борис Верховский99b77012019-11-02 15:09:14 -0400801 def flatten(list_of_lists):
Raymond Hettinger30c73622010-12-01 23:45:20 +0000802 "Flatten one level of nesting"
Борис Верховский99b77012019-11-02 15:09:14 -0400803 return chain.from_iterable(list_of_lists)
Georg Brandl116aa622007-08-15 14:28:22 +0000804
Raymond Hettinger30c73622010-12-01 23:45:20 +0000805 def repeatfunc(func, times=None, *args):
806 """Repeat calls to func with specified arguments.
Georg Brandl116aa622007-08-15 14:28:22 +0000807
Raymond Hettinger30c73622010-12-01 23:45:20 +0000808 Example: repeatfunc(random.random)
809 """
810 if times is None:
811 return starmap(func, repeat(args))
812 return starmap(func, repeat(args, times))
Georg Brandl116aa622007-08-15 14:28:22 +0000813
Raymond Hettinger44571da2013-05-05 19:53:41 -0700814 def grouper(iterable, n, fillvalue=None):
Raymond Hettinger9ae94732012-07-08 16:04:03 -0700815 "Collect data into fixed-length chunks or blocks"
Raymond Hettinger44571da2013-05-05 19:53:41 -0700816 # grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx"
Raymond Hettinger30c73622010-12-01 23:45:20 +0000817 args = [iter(iterable)] * n
818 return zip_longest(*args, fillvalue=fillvalue)
Georg Brandl116aa622007-08-15 14:28:22 +0000819
Raymond Hettinger30c73622010-12-01 23:45:20 +0000820 def roundrobin(*iterables):
821 "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
822 # Recipe credited to George Sakkis
Raymond Hettinger337cbba2017-11-21 00:23:34 -0800823 num_active = len(iterables)
Raymond Hettinger30c73622010-12-01 23:45:20 +0000824 nexts = cycle(iter(it).__next__ for it in iterables)
Raymond Hettinger337cbba2017-11-21 00:23:34 -0800825 while num_active:
Raymond Hettinger30c73622010-12-01 23:45:20 +0000826 try:
827 for next in nexts:
828 yield next()
829 except StopIteration:
Raymond Hettinger337cbba2017-11-21 00:23:34 -0800830 # Remove the iterator we just exhausted from the cycle.
831 num_active -= 1
832 nexts = cycle(islice(nexts, num_active))
Georg Brandl116aa622007-08-15 14:28:22 +0000833
Raymond Hettinger30c73622010-12-01 23:45:20 +0000834 def partition(pred, iterable):
Raymond Hettingerfc40b302020-11-29 10:47:22 -0800835 "Use a predicate to partition entries into false entries and true entries"
Raymond Hettinger30c73622010-12-01 23:45:20 +0000836 # partition(is_odd, range(10)) --> 0 2 4 6 8 and 1 3 5 7 9
837 t1, t2 = tee(iterable)
838 return filterfalse(pred, t1), filter(pred, t2)
Raymond Hettinger5ce0aa22010-12-01 10:49:19 +0000839
Raymond Hettinger30c73622010-12-01 23:45:20 +0000840 def powerset(iterable):
841 "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
842 s = list(iterable)
843 return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettinger08d01ee2010-08-07 05:36:53 +0000844
Raymond Hettinger30c73622010-12-01 23:45:20 +0000845 def unique_everseen(iterable, key=None):
846 "List unique elements, preserving order. Remember all elements ever seen."
847 # unique_everseen('AAAABBBCCDAABBB') --> A B C D
848 # unique_everseen('ABBCcAD', str.lower) --> A B C D
849 seen = set()
850 seen_add = seen.add
851 if key is None:
852 for element in filterfalse(seen.__contains__, iterable):
853 seen_add(element)
854 yield element
855 else:
856 for element in iterable:
857 k = key(element)
858 if k not in seen:
859 seen_add(k)
860 yield element
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000861
Raymond Hettinger30c73622010-12-01 23:45:20 +0000862 def unique_justseen(iterable, key=None):
863 "List unique elements, preserving order. Remember only the element just seen."
864 # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
865 # unique_justseen('ABBCcAD', str.lower) --> A B C A D
Jakub Molinskib4c7f392019-04-23 10:30:30 +0200866 return map(next, map(operator.itemgetter(1), groupby(iterable, key)))
Raymond Hettingerad9d96b2009-01-02 21:39:07 +0000867
Raymond Hettinger30c73622010-12-01 23:45:20 +0000868 def iter_except(func, exception, first=None):
869 """ Call a function repeatedly until an exception is raised.
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000870
Raymond Hettinger30c73622010-12-01 23:45:20 +0000871 Converts a call-until-exception interface to an iterator interface.
Andrew Kuchling1d7d5802013-06-20 21:17:41 -0400872 Like builtins.iter(func, sentinel) but uses an exception instead
Raymond Hettinger30c73622010-12-01 23:45:20 +0000873 of a sentinel to end the loop.
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000874
Raymond Hettinger30c73622010-12-01 23:45:20 +0000875 Examples:
876 iter_except(functools.partial(heappop, h), IndexError) # priority queue iterator
877 iter_except(d.popitem, KeyError) # non-blocking dict iterator
878 iter_except(d.popleft, IndexError) # non-blocking deque iterator
879 iter_except(q.get_nowait, Queue.Empty) # loop over a producer Queue
880 iter_except(s.pop, KeyError) # non-blocking set iterator
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000881
Raymond Hettinger30c73622010-12-01 23:45:20 +0000882 """
883 try:
884 if first is not None:
885 yield first() # For database APIs needing an initial cast to db.first()
Raymond Hettingera503f702016-03-13 00:12:31 -0800886 while True:
Raymond Hettinger30c73622010-12-01 23:45:20 +0000887 yield func()
888 except exception:
889 pass
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000890
Raymond Hettinger31b26f62014-04-02 03:16:42 -0700891 def first_true(iterable, default=False, pred=None):
892 """Returns the first true value in the iterable.
893
894 If no true value is found, returns *default*
895
896 If *pred* is not None, returns the first item
897 for which pred(item) is true.
898
899 """
900 # first_true([a,b,c], x) --> a or b or c or x
901 # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
902 return next(filter(pred, iterable), default)
903
Raymond Hettinger30c73622010-12-01 23:45:20 +0000904 def random_product(*args, repeat=1):
905 "Random selection from itertools.product(*args, **kwds)"
906 pools = [tuple(pool) for pool in args] * repeat
Raymond Hettingerfc40b302020-11-29 10:47:22 -0800907 return tuple(map(random.choice, pools))
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000908
Raymond Hettinger30c73622010-12-01 23:45:20 +0000909 def random_permutation(iterable, r=None):
910 "Random selection from itertools.permutations(iterable, r)"
911 pool = tuple(iterable)
912 r = len(pool) if r is None else r
913 return tuple(random.sample(pool, r))
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000914
Raymond Hettinger30c73622010-12-01 23:45:20 +0000915 def random_combination(iterable, r):
916 "Random selection from itertools.combinations(iterable, r)"
917 pool = tuple(iterable)
918 n = len(pool)
919 indices = sorted(random.sample(range(n), r))
920 return tuple(pool[i] for i in indices)
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000921
Raymond Hettinger30c73622010-12-01 23:45:20 +0000922 def random_combination_with_replacement(iterable, r):
923 "Random selection from itertools.combinations_with_replacement(iterable, r)"
924 pool = tuple(iterable)
925 n = len(pool)
Raymond Hettingerfc40b302020-11-29 10:47:22 -0800926 indices = sorted(random.choices(range(n), k=r))
Raymond Hettinger30c73622010-12-01 23:45:20 +0000927 return tuple(pool[i] for i in indices)
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000928
Raymond Hettingerd37258d2018-01-13 10:35:40 -0800929 def nth_combination(iterable, r, index):
Raymond Hettingerfc40b302020-11-29 10:47:22 -0800930 "Equivalent to list(combinations(iterable, r))[index]"
Raymond Hettingerd37258d2018-01-13 10:35:40 -0800931 pool = tuple(iterable)
932 n = len(pool)
933 if r < 0 or r > n:
934 raise ValueError
935 c = 1
936 k = min(r, n-r)
937 for i in range(1, k+1):
938 c = c * (n - k + i) // i
939 if index < 0:
940 index += c
941 if index < 0 or index >= c:
942 raise IndexError
943 result = []
944 while r:
945 c, n, r = c*r//n, n-1, r-1
946 while index >= c:
947 index -= c
948 c, n = c*(n-r)//n, n-1
949 result.append(pool[-1-n])
950 return tuple(result)
951