blob: b1513cd8b1b909d3febc5a59cbb76e29fc9ccfa1 [file] [log] [blame]
Georg Brandl116aa622007-08-15 14:28:22 +00001:mod:`itertools` --- Functions creating iterators for efficient looping
2=======================================================================
3
4.. module:: itertools
Raymond Hettinger30c73622010-12-01 23:45:20 +00005 :synopsis: Functions creating iterators for efficient looping.
Terry Jan Reedyfa089b92016-06-11 15:02:54 -04006
Georg Brandl116aa622007-08-15 14:28:22 +00007.. moduleauthor:: Raymond Hettinger <python@rcn.com>
8.. sectionauthor:: Raymond Hettinger <python@rcn.com>
9
Christian Heimesfe337bf2008-03-23 21:54:12 +000010.. testsetup::
11
Raymond Hettinger30c73622010-12-01 23:45:20 +000012 from itertools import *
Christian Heimesfe337bf2008-03-23 21:54:12 +000013
Terry Jan Reedyfa089b92016-06-11 15:02:54 -040014--------------
Christian Heimesfe337bf2008-03-23 21:54:12 +000015
Raymond Hettingerf76b9202009-02-17 20:00:59 +000016This module implements a number of :term:`iterator` building blocks inspired
17by constructs from APL, Haskell, and SML. Each has been recast in a form
18suitable for Python.
Georg Brandl116aa622007-08-15 14:28:22 +000019
20The module standardizes a core set of fast, memory efficient tools that are
Raymond Hettingerf76b9202009-02-17 20:00:59 +000021useful by themselves or in combination. Together, they form an "iterator
22algebra" making it possible to construct specialized tools succinctly and
23efficiently in pure Python.
Georg Brandl116aa622007-08-15 14:28:22 +000024
25For instance, SML provides a tabulation tool: ``tabulate(f)`` which produces a
Ezio Melottib6605992010-01-21 20:57:24 +000026sequence ``f(0), f(1), ...``. The same effect can be achieved in Python
Raymond Hettingera6c60372008-03-13 01:26:19 +000027by combining :func:`map` and :func:`count` to form ``map(f, count())``.
Georg Brandl116aa622007-08-15 14:28:22 +000028
Raymond Hettinger2c109ab2009-03-12 00:29:44 +000029These tools and their built-in counterparts also work well with the high-speed
30functions in the :mod:`operator` module. For example, the multiplication
31operator can be mapped across two vectors to form an efficient dot-product:
32``sum(map(operator.mul, vector1, vector2))``.
Georg Brandl116aa622007-08-15 14:28:22 +000033
Georg Brandl116aa622007-08-15 14:28:22 +000034
Raymond Hettinger6693d7a2017-12-15 13:17:55 -080035**Infinite iterators:**
Georg Brandl116aa622007-08-15 14:28:22 +000036
Raymond Hettinger5bfd8ce2009-04-10 19:02:36 +000037================== ================= ================================================= =========================================
38Iterator Arguments Results Example
39================== ================= ================================================= =========================================
40:func:`count` start, [step] start, start+step, start+2*step, ... ``count(10) --> 10 11 12 13 14 ...``
41:func:`cycle` p p0, p1, ... plast, p0, p1, ... ``cycle('ABCD') --> A B C D A B C D ...``
42:func:`repeat` elem [,n] elem, elem, elem, ... endlessly or up to n times ``repeat(10, 3) --> 10 10 10``
43================== ================= ================================================= =========================================
Georg Brandl116aa622007-08-15 14:28:22 +000044
Raymond Hettingerf76b9202009-02-17 20:00:59 +000045**Iterators terminating on the shortest input sequence:**
46
Benjamin Peterson2989f582014-01-16 10:10:13 -050047============================ ============================ ================================================= =============================================================
48Iterator Arguments Results Example
49============================ ============================ ================================================= =============================================================
50:func:`accumulate` p [,func] p0, p0+p1, p0+p1+p2, ... ``accumulate([1,2,3,4,5]) --> 1 3 6 10 15``
51:func:`chain` p, q, ... p0, p1, ... plast, q0, q1, ... ``chain('ABC', 'DEF') --> A B C D E F``
52:func:`chain.from_iterable` iterable p0, p1, ... plast, q0, q1, ... ``chain.from_iterable(['ABC', 'DEF']) --> A B C D E F``
53:func:`compress` data, selectors (d[0] if s[0]), (d[1] if s[1]), ... ``compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F``
54:func:`dropwhile` pred, seq seq[n], seq[n+1], starting when pred fails ``dropwhile(lambda x: x<5, [1,4,6,4,1]) --> 6 4 1``
55:func:`filterfalse` pred, seq elements of seq where pred(elem) is false ``filterfalse(lambda x: x%2, range(10)) --> 0 2 4 6 8``
Raymond Hettinger49392c62017-09-25 01:21:06 -070056:func:`groupby` iterable[, key] sub-iterators grouped by value of key(v)
Benjamin Peterson2989f582014-01-16 10:10:13 -050057:func:`islice` seq, [start,] stop [, step] elements from seq[start:stop:step] ``islice('ABCDEFG', 2, None) --> C D E F G``
58:func:`starmap` func, seq func(\*seq[0]), func(\*seq[1]), ... ``starmap(pow, [(2,5), (3,2), (10,3)]) --> 32 9 1000``
59:func:`takewhile` pred, seq seq[0], seq[1], until pred fails ``takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4``
60:func:`tee` it, n it1, it2, ... itn splits one iterator into n
61:func:`zip_longest` p, q, ... (p[0], q[0]), (p[1], q[1]), ... ``zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-``
62============================ ============================ ================================================= =============================================================
Raymond Hettingerf76b9202009-02-17 20:00:59 +000063
Raymond Hettinger6693d7a2017-12-15 13:17:55 -080064**Combinatoric iterators:**
Raymond Hettingerf76b9202009-02-17 20:00:59 +000065
Raymond Hettinger7f587cd2009-04-10 19:43:50 +000066============================================== ==================== =============================================================
67Iterator Arguments Results
68============================================== ==================== =============================================================
69:func:`product` p, q, ... [repeat=1] cartesian product, equivalent to a nested for-loop
70:func:`permutations` p[, r] r-length tuples, all possible orderings, no repeated elements
Raymond Hettinger36c3c022009-11-19 01:20:07 +000071:func:`combinations` p, r r-length tuples, in sorted order, no repeated elements
72:func:`combinations_with_replacement` p, r r-length tuples, in sorted order, with repeated elements
Raymond Hettinger7f587cd2009-04-10 19:43:50 +000073``product('ABCD', repeat=2)`` ``AA AB AC AD BA BB BC BD CA CB CC CD DA DB DC DD``
74``permutations('ABCD', 2)`` ``AB AC AD BA BC BD CA CB CD DA DB DC``
75``combinations('ABCD', 2)`` ``AB AC AD BC BD CD``
76``combinations_with_replacement('ABCD', 2)`` ``AA AB AC AD BB BC BD CC CD DD``
77============================================== ==================== =============================================================
Georg Brandl116aa622007-08-15 14:28:22 +000078
79
80.. _itertools-functions:
81
82Itertool functions
83------------------
84
85The following module functions all construct and return iterators. Some provide
86streams of infinite length, so they should only be accessed by functions or
87loops that truncate the stream.
88
Lisa Roach9718b592018-09-23 17:34:59 -070089.. function:: accumulate(iterable[, func, *, initial=None])
Raymond Hettingeradb81462010-12-01 22:50:36 +000090
Andrew Kuchling15b04eb2014-04-15 22:28:40 -040091 Make an iterator that returns accumulated sums, or accumulated
92 results of other binary functions (specified via the optional
Lisa Roach9718b592018-09-23 17:34:59 -070093 *func* argument).
94
95 If *func* is supplied, it should be a function
Andrew Kuchling15b04eb2014-04-15 22:28:40 -040096 of two arguments. Elements of the input *iterable* may be any type
97 that can be accepted as arguments to *func*. (For example, with
98 the default operation of addition, elements may be any addable
99 type including :class:`~decimal.Decimal` or
Lisa Roach9718b592018-09-23 17:34:59 -0700100 :class:`~fractions.Fraction`.)
101
102 Usually, the number of elements output matches the input iterable.
103 However, if the keyword argument *initial* is provided, the
104 accumulation leads off with the *initial* value so that the output
105 has one more element than the input iterable.
Raymond Hettingeradb81462010-12-01 22:50:36 +0000106
Raymond Hettinger672866d2016-05-28 00:17:54 -0700107 Roughly equivalent to::
Raymond Hettinger5d446132011-03-27 18:52:10 -0700108
Lisa Roach9718b592018-09-23 17:34:59 -0700109 def accumulate(iterable, func=operator.add, *, initial=None):
Raymond Hettingeradb81462010-12-01 22:50:36 +0000110 'Return running totals'
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000111 # accumulate([1,2,3,4,5]) --> 1 3 6 10 15
Lisa Roach9718b592018-09-23 17:34:59 -0700112 # accumulate([1,2,3,4,5], initial=100) --> 100 101 103 106 110 115
Raymond Hettinger5d446132011-03-27 18:52:10 -0700113 # accumulate([1,2,3,4,5], operator.mul) --> 1 2 6 24 120
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000114 it = iter(iterable)
Lisa Roach9718b592018-09-23 17:34:59 -0700115 total = initial
116 if initial is None:
117 try:
118 total = next(it)
119 except StopIteration:
120 return
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000121 yield total
122 for element in it:
Raymond Hettinger5d446132011-03-27 18:52:10 -0700123 total = func(total, element)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000124 yield total
Raymond Hettingeradb81462010-12-01 22:50:36 +0000125
Raymond Hettinger295c1d42011-04-21 11:09:28 -0700126 There are a number of uses for the *func* argument. It can be set to
127 :func:`min` for a running minimum, :func:`max` for a running maximum, or
128 :func:`operator.mul` for a running product. Amortization tables can be
129 built by accumulating interest and applying payments. First-order
Georg Brandl5d941342016-02-26 19:37:12 +0100130 `recurrence relations <https://en.wikipedia.org/wiki/Recurrence_relation>`_
Raymond Hettinger295c1d42011-04-21 11:09:28 -0700131 can be modeled by supplying the initial value in the iterable and using only
132 the accumulated total in *func* argument::
Raymond Hettinger5d446132011-03-27 18:52:10 -0700133
134 >>> data = [3, 4, 6, 2, 1, 9, 0, 7, 5, 8]
135 >>> list(accumulate(data, operator.mul)) # running product
136 [3, 12, 72, 144, 144, 1296, 0, 0, 0, 0]
137 >>> list(accumulate(data, max)) # running maximum
138 [3, 4, 6, 6, 6, 9, 9, 9, 9, 9]
139
140 # Amortize a 5% loan of 1000 with 4 annual payments of 90
141 >>> cashflows = [1000, -90, -90, -90, -90]
142 >>> list(accumulate(cashflows, lambda bal, pmt: bal*1.05 + pmt))
143 [1000, 960.0, 918.0, 873.9000000000001, 827.5950000000001]
144
Georg Brandl5d941342016-02-26 19:37:12 +0100145 # Chaotic recurrence relation https://en.wikipedia.org/wiki/Logistic_map
Raymond Hettinger295c1d42011-04-21 11:09:28 -0700146 >>> logistic_map = lambda x, _: r * x * (1 - x)
147 >>> r = 3.8
148 >>> x0 = 0.4
149 >>> inputs = repeat(x0, 36) # only the initial value is used
150 >>> [format(x, '.2f') for x in accumulate(inputs, logistic_map)]
151 ['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 +0200152 '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 -0700153 '0.93', '0.25', '0.71', '0.79', '0.63', '0.88', '0.39', '0.91', '0.32',
154 '0.83', '0.54', '0.95', '0.20', '0.60', '0.91', '0.30', '0.80', '0.60']
155
Raymond Hettinger64801682013-10-12 16:04:17 -0700156 See :func:`functools.reduce` for a similar function that returns only the
157 final accumulated value.
158
Raymond Hettingeradb81462010-12-01 22:50:36 +0000159 .. versionadded:: 3.2
Georg Brandl116aa622007-08-15 14:28:22 +0000160
Raymond Hettinger5d446132011-03-27 18:52:10 -0700161 .. versionchanged:: 3.3
162 Added the optional *func* parameter.
163
Lisa Roach9718b592018-09-23 17:34:59 -0700164 .. versionchanged:: 3.8
165 Added the optional *initial* parameter.
166
Georg Brandl116aa622007-08-15 14:28:22 +0000167.. function:: chain(*iterables)
168
Raymond Hettinger30c73622010-12-01 23:45:20 +0000169 Make an iterator that returns elements from the first iterable until it is
170 exhausted, then proceeds to the next iterable, until all of the iterables are
171 exhausted. Used for treating consecutive sequences as a single sequence.
Raymond Hettinger672866d2016-05-28 00:17:54 -0700172 Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000173
Raymond Hettinger30c73622010-12-01 23:45:20 +0000174 def chain(*iterables):
175 # chain('ABC', 'DEF') --> A B C D E F
176 for it in iterables:
177 for element in it:
178 yield element
Georg Brandl116aa622007-08-15 14:28:22 +0000179
180
Georg Brandl933b9742010-07-29 14:36:11 +0000181.. classmethod:: chain.from_iterable(iterable)
Christian Heimes70e7ea22008-02-28 20:02:27 +0000182
Raymond Hettinger30c73622010-12-01 23:45:20 +0000183 Alternate constructor for :func:`chain`. Gets chained inputs from a
Raymond Hettinger1e21ebc2013-09-09 01:54:27 -0500184 single iterable argument that is evaluated lazily. Roughly equivalent to::
Christian Heimes70e7ea22008-02-28 20:02:27 +0000185
Raymond Hettinger30c73622010-12-01 23:45:20 +0000186 def from_iterable(iterables):
187 # chain.from_iterable(['ABC', 'DEF']) --> A B C D E F
188 for it in iterables:
189 for element in it:
190 yield element
Christian Heimes70e7ea22008-02-28 20:02:27 +0000191
Christian Heimes78644762008-03-04 23:39:23 +0000192
Christian Heimes836baa52008-02-26 08:18:30 +0000193.. function:: combinations(iterable, r)
194
Raymond Hettinger30c73622010-12-01 23:45:20 +0000195 Return *r* length subsequences of elements from the input *iterable*.
Christian Heimes836baa52008-02-26 08:18:30 +0000196
Raymond Hettinger30c73622010-12-01 23:45:20 +0000197 Combinations are emitted in lexicographic sort order. So, if the
198 input *iterable* is sorted, the combination tuples will be produced
199 in sorted order.
Christian Heimes836baa52008-02-26 08:18:30 +0000200
Raymond Hettinger30c73622010-12-01 23:45:20 +0000201 Elements are treated as unique based on their position, not on their
202 value. So if the input elements are unique, there will be no repeat
203 values in each combination.
Christian Heimes836baa52008-02-26 08:18:30 +0000204
Raymond Hettinger672866d2016-05-28 00:17:54 -0700205 Roughly equivalent to::
Christian Heimes836baa52008-02-26 08:18:30 +0000206
207 def combinations(iterable, r):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000208 # combinations('ABCD', 2) --> AB AC AD BC BD CD
209 # combinations(range(4), 3) --> 012 013 023 123
Christian Heimes836baa52008-02-26 08:18:30 +0000210 pool = tuple(iterable)
Christian Heimes380f7f22008-02-28 11:19:05 +0000211 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000212 if r > n:
213 return
214 indices = list(range(r))
Christian Heimesb558a2e2008-03-02 22:46:37 +0000215 yield tuple(pool[i] for i in indices)
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000216 while True:
Christian Heimes380f7f22008-02-28 11:19:05 +0000217 for i in reversed(range(r)):
Christian Heimesb558a2e2008-03-02 22:46:37 +0000218 if indices[i] != i + n - r:
Christian Heimes836baa52008-02-26 08:18:30 +0000219 break
Christian Heimes380f7f22008-02-28 11:19:05 +0000220 else:
221 return
Christian Heimesb558a2e2008-03-02 22:46:37 +0000222 indices[i] += 1
Christian Heimes380f7f22008-02-28 11:19:05 +0000223 for j in range(i+1, r):
Christian Heimesb558a2e2008-03-02 22:46:37 +0000224 indices[j] = indices[j-1] + 1
225 yield tuple(pool[i] for i in indices)
Christian Heimes836baa52008-02-26 08:18:30 +0000226
Raymond Hettinger30c73622010-12-01 23:45:20 +0000227 The code for :func:`combinations` can be also expressed as a subsequence
228 of :func:`permutations` after filtering entries where the elements are not
229 in sorted order (according to their position in the input pool)::
Christian Heimes78644762008-03-04 23:39:23 +0000230
231 def combinations(iterable, r):
232 pool = tuple(iterable)
233 n = len(pool)
234 for indices in permutations(range(n), r):
235 if sorted(indices) == list(indices):
236 yield tuple(pool[i] for i in indices)
237
Raymond Hettinger30c73622010-12-01 23:45:20 +0000238 The number of items returned is ``n! / r! / (n-r)!`` when ``0 <= r <= n``
239 or zero when ``r > n``.
Christian Heimes836baa52008-02-26 08:18:30 +0000240
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000241.. function:: combinations_with_replacement(iterable, r)
242
Raymond Hettinger30c73622010-12-01 23:45:20 +0000243 Return *r* length subsequences of elements from the input *iterable*
244 allowing individual elements to be repeated more than once.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000245
Raymond Hettinger30c73622010-12-01 23:45:20 +0000246 Combinations are emitted in lexicographic sort order. So, if the
247 input *iterable* is sorted, the combination tuples will be produced
248 in sorted order.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000249
Raymond Hettinger30c73622010-12-01 23:45:20 +0000250 Elements are treated as unique based on their position, not on their
251 value. So if the input elements are unique, the generated combinations
252 will also be unique.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000253
Raymond Hettinger672866d2016-05-28 00:17:54 -0700254 Roughly equivalent to::
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000255
256 def combinations_with_replacement(iterable, r):
257 # combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC
258 pool = tuple(iterable)
259 n = len(pool)
260 if not n and r:
261 return
262 indices = [0] * r
263 yield tuple(pool[i] for i in indices)
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000264 while True:
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000265 for i in reversed(range(r)):
266 if indices[i] != n - 1:
267 break
268 else:
269 return
270 indices[i:] = [indices[i] + 1] * (r - i)
271 yield tuple(pool[i] for i in indices)
272
Raymond Hettinger30c73622010-12-01 23:45:20 +0000273 The code for :func:`combinations_with_replacement` can be also expressed as
274 a subsequence of :func:`product` after filtering entries where the elements
275 are not in sorted order (according to their position in the input pool)::
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000276
277 def combinations_with_replacement(iterable, r):
278 pool = tuple(iterable)
279 n = len(pool)
280 for indices in product(range(n), repeat=r):
281 if sorted(indices) == list(indices):
282 yield tuple(pool[i] for i in indices)
283
Raymond Hettinger30c73622010-12-01 23:45:20 +0000284 The number of items returned is ``(n+r-1)! / r! / (n-1)!`` when ``n > 0``.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000285
Raymond Hettinger30c73622010-12-01 23:45:20 +0000286 .. versionadded:: 3.1
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000287
Georg Brandl67b21b72010-08-17 15:07:14 +0000288
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000289.. function:: compress(data, selectors)
290
Raymond Hettinger30c73622010-12-01 23:45:20 +0000291 Make an iterator that filters elements from *data* returning only those that
292 have a corresponding element in *selectors* that evaluates to ``True``.
293 Stops when either the *data* or *selectors* iterables has been exhausted.
Raymond Hettinger672866d2016-05-28 00:17:54 -0700294 Roughly equivalent to::
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000295
Raymond Hettinger30c73622010-12-01 23:45:20 +0000296 def compress(data, selectors):
297 # compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F
298 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000299
Raymond Hettinger30c73622010-12-01 23:45:20 +0000300 .. versionadded:: 3.1
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000301
302
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000303.. function:: count(start=0, step=1)
Georg Brandl116aa622007-08-15 14:28:22 +0000304
Andrew Kuchlingedb42602013-06-21 08:00:58 -0400305 Make an iterator that returns evenly spaced values starting with number *start*. Often
Raymond Hettinger30c73622010-12-01 23:45:20 +0000306 used as an argument to :func:`map` to generate consecutive data points.
Raymond Hettinger672866d2016-05-28 00:17:54 -0700307 Also, used with :func:`zip` to add sequence numbers. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000308
Raymond Hettinger30c73622010-12-01 23:45:20 +0000309 def count(start=0, step=1):
310 # count(10) --> 10 11 12 13 14 ...
Georg Brandl37a80dc2011-01-13 07:31:18 +0000311 # count(2.5, 0.5) -> 2.5 3.0 3.5 ...
Raymond Hettinger30c73622010-12-01 23:45:20 +0000312 n = start
313 while True:
314 yield n
315 n += step
Georg Brandl116aa622007-08-15 14:28:22 +0000316
Raymond Hettinger30c73622010-12-01 23:45:20 +0000317 When counting with floating point numbers, better accuracy can sometimes be
318 achieved by substituting multiplicative code such as: ``(start + step * i
319 for i in count())``.
Raymond Hettinger5bc472a2009-06-17 01:40:52 +0000320
Raymond Hettinger30c73622010-12-01 23:45:20 +0000321 .. versionchanged:: 3.1
322 Added *step* argument and allowed non-integer arguments.
Georg Brandl116aa622007-08-15 14:28:22 +0000323
324.. function:: cycle(iterable)
325
Raymond Hettinger30c73622010-12-01 23:45:20 +0000326 Make an iterator returning elements from the iterable and saving a copy of each.
327 When the iterable is exhausted, return elements from the saved copy. Repeats
Raymond Hettinger672866d2016-05-28 00:17:54 -0700328 indefinitely. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000329
Raymond Hettinger30c73622010-12-01 23:45:20 +0000330 def cycle(iterable):
331 # cycle('ABCD') --> A B C D A B C D A B C D ...
332 saved = []
333 for element in iterable:
334 yield element
335 saved.append(element)
336 while saved:
337 for element in saved:
Georg Brandl116aa622007-08-15 14:28:22 +0000338 yield element
339
Raymond Hettinger30c73622010-12-01 23:45:20 +0000340 Note, this member of the toolkit may require significant auxiliary storage
341 (depending on the length of the iterable).
Georg Brandl116aa622007-08-15 14:28:22 +0000342
343
344.. function:: dropwhile(predicate, iterable)
345
Raymond Hettinger30c73622010-12-01 23:45:20 +0000346 Make an iterator that drops elements from the iterable as long as the predicate
347 is true; afterwards, returns every element. Note, the iterator does not produce
348 *any* output until the predicate first becomes false, so it may have a lengthy
Raymond Hettinger672866d2016-05-28 00:17:54 -0700349 start-up time. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000350
Raymond Hettinger30c73622010-12-01 23:45:20 +0000351 def dropwhile(predicate, iterable):
352 # dropwhile(lambda x: x<5, [1,4,6,4,1]) --> 6 4 1
353 iterable = iter(iterable)
354 for x in iterable:
355 if not predicate(x):
356 yield x
357 break
358 for x in iterable:
359 yield x
Georg Brandl116aa622007-08-15 14:28:22 +0000360
Raymond Hettinger749761e2009-01-27 04:42:48 +0000361.. function:: filterfalse(predicate, iterable)
362
Raymond Hettinger30c73622010-12-01 23:45:20 +0000363 Make an iterator that filters elements from iterable returning only those for
364 which the predicate is ``False``. If *predicate* is ``None``, return the items
Raymond Hettinger672866d2016-05-28 00:17:54 -0700365 that are false. Roughly equivalent to::
Raymond Hettinger749761e2009-01-27 04:42:48 +0000366
Raymond Hettinger30c73622010-12-01 23:45:20 +0000367 def filterfalse(predicate, iterable):
368 # filterfalse(lambda x: x%2, range(10)) --> 0 2 4 6 8
369 if predicate is None:
370 predicate = bool
371 for x in iterable:
372 if not predicate(x):
373 yield x
Raymond Hettinger749761e2009-01-27 04:42:48 +0000374
Georg Brandl116aa622007-08-15 14:28:22 +0000375
Georg Brandl3dd33882009-06-01 17:35:27 +0000376.. function:: groupby(iterable, key=None)
Georg Brandl116aa622007-08-15 14:28:22 +0000377
Raymond Hettinger30c73622010-12-01 23:45:20 +0000378 Make an iterator that returns consecutive keys and groups from the *iterable*.
379 The *key* is a function computing a key value for each element. If not
380 specified or is ``None``, *key* defaults to an identity function and returns
381 the element unchanged. Generally, the iterable needs to already be sorted on
382 the same key function.
Georg Brandl116aa622007-08-15 14:28:22 +0000383
Raymond Hettinger30c73622010-12-01 23:45:20 +0000384 The operation of :func:`groupby` is similar to the ``uniq`` filter in Unix. It
385 generates a break or new group every time the value of the key function changes
386 (which is why it is usually necessary to have sorted the data using the same key
387 function). That behavior differs from SQL's GROUP BY which aggregates common
388 elements regardless of their input order.
Georg Brandl116aa622007-08-15 14:28:22 +0000389
Raymond Hettinger30c73622010-12-01 23:45:20 +0000390 The returned group is itself an iterator that shares the underlying iterable
391 with :func:`groupby`. Because the source is shared, when the :func:`groupby`
392 object is advanced, the previous group is no longer visible. So, if that data
393 is needed later, it should be stored as a list::
Georg Brandl116aa622007-08-15 14:28:22 +0000394
Raymond Hettinger30c73622010-12-01 23:45:20 +0000395 groups = []
396 uniquekeys = []
397 data = sorted(data, key=keyfunc)
398 for k, g in groupby(data, keyfunc):
399 groups.append(list(g)) # Store group iterator as a list
400 uniquekeys.append(k)
Georg Brandl116aa622007-08-15 14:28:22 +0000401
Raymond Hettinger672866d2016-05-28 00:17:54 -0700402 :func:`groupby` is roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000403
Raymond Hettinger30c73622010-12-01 23:45:20 +0000404 class groupby:
405 # [k for k, g in groupby('AAAABBBCCDAABBB')] --> A B C D A B
406 # [list(g) for k, g in groupby('AAAABBBCCD')] --> AAAA BBB CC D
407 def __init__(self, iterable, key=None):
408 if key is None:
409 key = lambda x: x
410 self.keyfunc = key
411 self.it = iter(iterable)
412 self.tgtkey = self.currkey = self.currvalue = object()
413 def __iter__(self):
414 return self
415 def __next__(self):
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300416 self.id = object()
Raymond Hettinger30c73622010-12-01 23:45:20 +0000417 while self.currkey == self.tgtkey:
418 self.currvalue = next(self.it) # Exit on StopIteration
419 self.currkey = self.keyfunc(self.currvalue)
420 self.tgtkey = self.currkey
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300421 return (self.currkey, self._grouper(self.tgtkey, self.id))
422 def _grouper(self, tgtkey, id):
423 while self.id is id and self.currkey == tgtkey:
Raymond Hettinger30c73622010-12-01 23:45:20 +0000424 yield self.currvalue
Raymond Hettinger828d9322014-11-22 21:56:23 -0800425 try:
426 self.currvalue = next(self.it)
427 except StopIteration:
428 return
Raymond Hettinger30c73622010-12-01 23:45:20 +0000429 self.currkey = self.keyfunc(self.currvalue)
Georg Brandl116aa622007-08-15 14:28:22 +0000430
Georg Brandl116aa622007-08-15 14:28:22 +0000431
Ezio Melottie0add762012-09-14 06:32:35 +0300432.. function:: islice(iterable, stop)
433 islice(iterable, start, stop[, step])
Georg Brandl116aa622007-08-15 14:28:22 +0000434
Raymond Hettinger30c73622010-12-01 23:45:20 +0000435 Make an iterator that returns selected elements from the iterable. If *start* is
436 non-zero, then elements from the iterable are skipped until start is reached.
437 Afterward, elements are returned consecutively unless *step* is set higher than
438 one which results in items being skipped. If *stop* is ``None``, then iteration
439 continues until the iterator is exhausted, if at all; otherwise, it stops at the
440 specified position. Unlike regular slicing, :func:`islice` does not support
441 negative values for *start*, *stop*, or *step*. Can be used to extract related
442 fields from data where the internal structure has been flattened (for example, a
Raymond Hettinger672866d2016-05-28 00:17:54 -0700443 multi-line report may list a name field on every third line). Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000444
Raymond Hettinger30c73622010-12-01 23:45:20 +0000445 def islice(iterable, *args):
446 # islice('ABCDEFG', 2) --> A B
447 # islice('ABCDEFG', 2, 4) --> C D
448 # islice('ABCDEFG', 2, None) --> C D E F G
449 # islice('ABCDEFG', 0, None, 2) --> A C E G
450 s = slice(*args)
Cheryl Sabellada1734c2018-03-26 21:29:33 -0400451 start, stop, step = s.start or 0, s.stop or sys.maxsize, s.step or 1
452 it = iter(range(start, stop, step))
Raymond Hettinger828d9322014-11-22 21:56:23 -0800453 try:
454 nexti = next(it)
455 except StopIteration:
Cheryl Sabellada1734c2018-03-26 21:29:33 -0400456 # Consume *iterable* up to the *start* position.
457 for i, element in zip(range(start), iterable):
458 pass
Raymond Hettinger828d9322014-11-22 21:56:23 -0800459 return
Cheryl Sabellada1734c2018-03-26 21:29:33 -0400460 try:
461 for i, element in enumerate(iterable):
462 if i == nexti:
463 yield element
464 nexti = next(it)
465 except StopIteration:
466 # Consume to *stop*.
467 for i, element in zip(range(i + 1, stop), iterable):
468 pass
Georg Brandl116aa622007-08-15 14:28:22 +0000469
Raymond Hettinger30c73622010-12-01 23:45:20 +0000470 If *start* is ``None``, then iteration starts at zero. If *step* is ``None``,
471 then the step defaults to one.
Georg Brandl116aa622007-08-15 14:28:22 +0000472
Georg Brandl116aa622007-08-15 14:28:22 +0000473
Georg Brandl3dd33882009-06-01 17:35:27 +0000474.. function:: permutations(iterable, r=None)
Christian Heimes70e7ea22008-02-28 20:02:27 +0000475
Raymond Hettinger30c73622010-12-01 23:45:20 +0000476 Return successive *r* length permutations of elements in the *iterable*.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000477
Raymond Hettinger30c73622010-12-01 23:45:20 +0000478 If *r* is not specified or is ``None``, then *r* defaults to the length
479 of the *iterable* and all possible full-length permutations
480 are generated.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000481
Raymond Hettinger30c73622010-12-01 23:45:20 +0000482 Permutations are emitted in lexicographic sort order. So, if the
483 input *iterable* is sorted, the permutation tuples will be produced
484 in sorted order.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000485
Raymond Hettinger30c73622010-12-01 23:45:20 +0000486 Elements are treated as unique based on their position, not on their
487 value. So if the input elements are unique, there will be no repeat
488 values in each permutation.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000489
Raymond Hettinger672866d2016-05-28 00:17:54 -0700490 Roughly equivalent to::
Christian Heimesb558a2e2008-03-02 22:46:37 +0000491
492 def permutations(iterable, r=None):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000493 # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
494 # permutations(range(3)) --> 012 021 102 120 201 210
Christian Heimesb558a2e2008-03-02 22:46:37 +0000495 pool = tuple(iterable)
496 n = len(pool)
497 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000498 if r > n:
499 return
500 indices = list(range(n))
Sandro Tosi73866622011-12-25 17:25:45 +0100501 cycles = list(range(n, n-r, -1))
Christian Heimesb558a2e2008-03-02 22:46:37 +0000502 yield tuple(pool[i] for i in indices[:r])
503 while n:
504 for i in reversed(range(r)):
505 cycles[i] -= 1
506 if cycles[i] == 0:
507 indices[i:] = indices[i+1:] + indices[i:i+1]
508 cycles[i] = n - i
509 else:
510 j = cycles[i]
511 indices[i], indices[-j] = indices[-j], indices[i]
512 yield tuple(pool[i] for i in indices[:r])
513 break
514 else:
515 return
Christian Heimes70e7ea22008-02-28 20:02:27 +0000516
Raymond Hettinger30c73622010-12-01 23:45:20 +0000517 The code for :func:`permutations` can be also expressed as a subsequence of
518 :func:`product`, filtered to exclude entries with repeated elements (those
519 from the same position in the input pool)::
Christian Heimes78644762008-03-04 23:39:23 +0000520
521 def permutations(iterable, r=None):
522 pool = tuple(iterable)
523 n = len(pool)
524 r = n if r is None else r
525 for indices in product(range(n), repeat=r):
526 if len(set(indices)) == r:
527 yield tuple(pool[i] for i in indices)
528
Raymond Hettinger30c73622010-12-01 23:45:20 +0000529 The number of items returned is ``n! / (n-r)!`` when ``0 <= r <= n``
530 or zero when ``r > n``.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000531
Georg Brandl3dd33882009-06-01 17:35:27 +0000532.. function:: product(*iterables, repeat=1)
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000533
Raymond Hettinger30c73622010-12-01 23:45:20 +0000534 Cartesian product of input iterables.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000535
Raymond Hettinger672866d2016-05-28 00:17:54 -0700536 Roughly equivalent to nested for-loops in a generator expression. For example,
Raymond Hettinger30c73622010-12-01 23:45:20 +0000537 ``product(A, B)`` returns the same as ``((x,y) for x in A for y in B)``.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000538
Raymond Hettinger30c73622010-12-01 23:45:20 +0000539 The nested loops cycle like an odometer with the rightmost element advancing
540 on every iteration. This pattern creates a lexicographic ordering so that if
541 the input's iterables are sorted, the product tuples are emitted in sorted
542 order.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000543
Raymond Hettinger30c73622010-12-01 23:45:20 +0000544 To compute the product of an iterable with itself, specify the number of
545 repetitions with the optional *repeat* keyword argument. For example,
546 ``product(A, repeat=4)`` means the same as ``product(A, A, A, A)``.
Christian Heimes9e7f1d22008-02-28 12:27:11 +0000547
Raymond Hettinger672866d2016-05-28 00:17:54 -0700548 This function is roughly equivalent to the following code, except that the
Raymond Hettinger30c73622010-12-01 23:45:20 +0000549 actual implementation does not build up intermediate results in memory::
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000550
Raymond Hettinger30c73622010-12-01 23:45:20 +0000551 def product(*args, repeat=1):
552 # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy
553 # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111
554 pools = [tuple(pool) for pool in args] * repeat
555 result = [[]]
556 for pool in pools:
557 result = [x+[y] for x in result for y in pool]
558 for prod in result:
559 yield tuple(prod)
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000560
561
Raymond Hettingerd75ad442009-06-01 19:16:52 +0000562.. function:: repeat(object[, times])
Georg Brandl116aa622007-08-15 14:28:22 +0000563
Raymond Hettinger30c73622010-12-01 23:45:20 +0000564 Make an iterator that returns *object* over and over again. Runs indefinitely
565 unless the *times* argument is specified. Used as argument to :func:`map` for
566 invariant parameters to the called function. Also used with :func:`zip` to
Raymond Hettinger672866d2016-05-28 00:17:54 -0700567 create an invariant part of a tuple record.
568
569 Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000570
Raymond Hettinger30c73622010-12-01 23:45:20 +0000571 def repeat(object, times=None):
572 # repeat(10, 3) --> 10 10 10
573 if times is None:
574 while True:
575 yield object
576 else:
577 for i in range(times):
578 yield object
Georg Brandl116aa622007-08-15 14:28:22 +0000579
Raymond Hettingerfc3ba6b2012-02-01 09:07:40 -0800580 A common use for *repeat* is to supply a stream of constant values to *map*
581 or *zip*::
582
583 >>> list(map(pow, range(10), repeat(2)))
584 [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
Georg Brandl116aa622007-08-15 14:28:22 +0000585
586.. function:: starmap(function, iterable)
587
Raymond Hettinger30c73622010-12-01 23:45:20 +0000588 Make an iterator that computes the function using arguments obtained from
589 the iterable. Used instead of :func:`map` when argument parameters are already
590 grouped in tuples from a single iterable (the data has been "pre-zipped"). The
591 difference between :func:`map` and :func:`starmap` parallels the distinction
Raymond Hettinger672866d2016-05-28 00:17:54 -0700592 between ``function(a,b)`` and ``function(*c)``. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000593
Raymond Hettinger30c73622010-12-01 23:45:20 +0000594 def starmap(function, iterable):
595 # starmap(pow, [(2,5), (3,2), (10,3)]) --> 32 9 1000
596 for args in iterable:
597 yield function(*args)
Christian Heimes679db4a2008-01-18 09:56:22 +0000598
Georg Brandl116aa622007-08-15 14:28:22 +0000599
600.. function:: takewhile(predicate, iterable)
601
Raymond Hettinger30c73622010-12-01 23:45:20 +0000602 Make an iterator that returns elements from the iterable as long as the
Raymond Hettinger672866d2016-05-28 00:17:54 -0700603 predicate is true. Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000604
Raymond Hettinger30c73622010-12-01 23:45:20 +0000605 def takewhile(predicate, iterable):
606 # takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4
607 for x in iterable:
608 if predicate(x):
609 yield x
610 else:
611 break
Georg Brandl116aa622007-08-15 14:28:22 +0000612
613
Georg Brandl3dd33882009-06-01 17:35:27 +0000614.. function:: tee(iterable, n=2)
Georg Brandl116aa622007-08-15 14:28:22 +0000615
Raymond Hettingerab425aa2016-04-26 00:10:00 -0700616 Return *n* independent iterators from a single iterable.
617
618 The following Python code helps explain what *tee* does (although the actual
Serhiy Storchaka4ecfa452016-05-16 09:31:54 +0300619 implementation is more complex and uses only a single underlying
Raymond Hettinger672866d2016-05-28 00:17:54 -0700620 :abbr:`FIFO (first-in, first-out)` queue).
621
622 Roughly equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000623
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000624 def tee(iterable, n=2):
625 it = iter(iterable)
626 deques = [collections.deque() for i in range(n)]
627 def gen(mydeque):
628 while True:
629 if not mydeque: # when the local deque is empty
Raymond Hettinger828d9322014-11-22 21:56:23 -0800630 try:
631 newval = next(it) # fetch a new value and
632 except StopIteration:
633 return
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000634 for d in deques: # load it to all the deques
635 d.append(newval)
636 yield mydeque.popleft()
637 return tuple(gen(d) for d in deques)
Georg Brandl116aa622007-08-15 14:28:22 +0000638
Raymond Hettinger30c73622010-12-01 23:45:20 +0000639 Once :func:`tee` has made a split, the original *iterable* should not be
640 used anywhere else; otherwise, the *iterable* could get advanced without
641 the tee objects being informed.
Georg Brandl116aa622007-08-15 14:28:22 +0000642
Raymond Hettinger30c73622010-12-01 23:45:20 +0000643 This itertool may require significant auxiliary storage (depending on how
644 much temporary data needs to be stored). In general, if one iterator uses
645 most or all of the data before another iterator starts, it is faster to use
646 :func:`list` instead of :func:`tee`.
Georg Brandl116aa622007-08-15 14:28:22 +0000647
Georg Brandl116aa622007-08-15 14:28:22 +0000648
Georg Brandl3dd33882009-06-01 17:35:27 +0000649.. function:: zip_longest(*iterables, fillvalue=None)
Raymond Hettinger749761e2009-01-27 04:42:48 +0000650
Raymond Hettinger30c73622010-12-01 23:45:20 +0000651 Make an iterator that aggregates elements from each of the iterables. If the
652 iterables are of uneven length, missing values are filled-in with *fillvalue*.
Raymond Hettinger672866d2016-05-28 00:17:54 -0700653 Iteration continues until the longest iterable is exhausted. Roughly equivalent to::
Raymond Hettinger749761e2009-01-27 04:42:48 +0000654
Raymond Hettinger3147b042017-09-07 14:01:44 -0700655 def zip_longest(*args, fillvalue=None):
Raymond Hettinger30c73622010-12-01 23:45:20 +0000656 # zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-
Raymond Hettinger3147b042017-09-07 14:01:44 -0700657 iterators = [iter(it) for it in args]
658 num_active = len(iterators)
659 if not num_active:
660 return
661 while True:
662 values = []
663 for i, it in enumerate(iterators):
664 try:
665 value = next(it)
666 except StopIteration:
667 num_active -= 1
668 if not num_active:
669 return
670 iterators[i] = repeat(fillvalue)
671 value = fillvalue
672 values.append(value)
673 yield tuple(values)
Raymond Hettinger749761e2009-01-27 04:42:48 +0000674
Raymond Hettinger30c73622010-12-01 23:45:20 +0000675 If one of the iterables is potentially infinite, then the :func:`zip_longest`
676 function should be wrapped with something that limits the number of calls
677 (for example :func:`islice` or :func:`takewhile`). If not specified,
678 *fillvalue* defaults to ``None``.
Raymond Hettinger749761e2009-01-27 04:42:48 +0000679
680
Georg Brandl116aa622007-08-15 14:28:22 +0000681.. _itertools-recipes:
682
Raymond Hettinger1fa76822010-12-06 23:31:36 +0000683Itertools Recipes
684-----------------
Georg Brandl116aa622007-08-15 14:28:22 +0000685
686This section shows recipes for creating an extended toolset using the existing
687itertools as building blocks.
688
689The extended tools offer the same high performance as the underlying toolset.
690The superior memory performance is kept by processing elements one at a time
691rather than bringing the whole iterable into memory all at once. Code volume is
692kept small by linking the tools together in a functional style which helps
693eliminate temporary variables. High speed is retained by preferring
Georg Brandl9afde1c2007-11-01 20:32:30 +0000694"vectorized" building blocks over the use of for-loops and :term:`generator`\s
Christian Heimesfe337bf2008-03-23 21:54:12 +0000695which incur interpreter overhead.
696
697.. testcode::
Georg Brandl116aa622007-08-15 14:28:22 +0000698
Raymond Hettinger30c73622010-12-01 23:45:20 +0000699 def take(n, iterable):
700 "Return first n items of the iterable as a list"
701 return list(islice(iterable, n))
Georg Brandl116aa622007-08-15 14:28:22 +0000702
Raymond Hettinger9265dd72018-04-08 08:44:20 -0700703 def prepend(value, iterator):
704 "Prepend a single value in front of an iterator"
705 # prepend(1, [2, 3, 4]) -> 1 2 3 4
706 return chain([value], iterator)
707
Raymond Hettinger30c73622010-12-01 23:45:20 +0000708 def tabulate(function, start=0):
709 "Return function(0), function(1), ..."
710 return map(function, count(start))
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000711
Raymond Hettingerdfe098d2014-05-25 22:03:56 -0700712 def tail(n, iterable):
713 "Return an iterator over the last n items"
714 # tail(3, 'ABCDEFG') --> E F G
715 return iter(collections.deque(iterable, maxlen=n))
716
Cheryl Sabellada1734c2018-03-26 21:29:33 -0400717 def consume(iterator, n=None):
718 "Advance the iterator n-steps ahead. If n is None, consume entirely."
Raymond Hettinger30c73622010-12-01 23:45:20 +0000719 # Use functions that consume iterators at C speed.
720 if n is None:
721 # feed the entire iterator into a zero-length deque
722 collections.deque(iterator, maxlen=0)
723 else:
724 # advance to the empty slice starting at position n
725 next(islice(iterator, n, n), None)
Raymond Hettingerfa007962009-03-09 11:55:25 +0000726
Raymond Hettinger30c73622010-12-01 23:45:20 +0000727 def nth(iterable, n, default=None):
728 "Returns the nth item or a default value"
729 return next(islice(iterable, n, None), default)
Georg Brandl116aa622007-08-15 14:28:22 +0000730
Raymond Hettingere525ee32016-03-06 18:11:38 -0800731 def all_equal(iterable):
732 "Returns True if all the elements are equal to each other"
733 g = groupby(iterable)
734 return next(g, True) and not next(g, False)
735
Raymond Hettinger30c73622010-12-01 23:45:20 +0000736 def quantify(iterable, pred=bool):
737 "Count how many times the predicate is true"
738 return sum(map(pred, iterable))
Georg Brandl116aa622007-08-15 14:28:22 +0000739
Raymond Hettinger30c73622010-12-01 23:45:20 +0000740 def padnone(iterable):
741 """Returns the sequence elements and then returns None indefinitely.
Georg Brandl116aa622007-08-15 14:28:22 +0000742
Raymond Hettinger30c73622010-12-01 23:45:20 +0000743 Useful for emulating the behavior of the built-in map() function.
744 """
745 return chain(iterable, repeat(None))
Georg Brandl116aa622007-08-15 14:28:22 +0000746
Raymond Hettinger30c73622010-12-01 23:45:20 +0000747 def ncycles(iterable, n):
748 "Returns the sequence elements n times"
749 return chain.from_iterable(repeat(tuple(iterable), n))
Georg Brandl116aa622007-08-15 14:28:22 +0000750
Raymond Hettinger30c73622010-12-01 23:45:20 +0000751 def dotproduct(vec1, vec2):
752 return sum(map(operator.mul, vec1, vec2))
Georg Brandl116aa622007-08-15 14:28:22 +0000753
Raymond Hettinger30c73622010-12-01 23:45:20 +0000754 def flatten(listOfLists):
755 "Flatten one level of nesting"
756 return chain.from_iterable(listOfLists)
Georg Brandl116aa622007-08-15 14:28:22 +0000757
Raymond Hettinger30c73622010-12-01 23:45:20 +0000758 def repeatfunc(func, times=None, *args):
759 """Repeat calls to func with specified arguments.
Georg Brandl116aa622007-08-15 14:28:22 +0000760
Raymond Hettinger30c73622010-12-01 23:45:20 +0000761 Example: repeatfunc(random.random)
762 """
763 if times is None:
764 return starmap(func, repeat(args))
765 return starmap(func, repeat(args, times))
Georg Brandl116aa622007-08-15 14:28:22 +0000766
Raymond Hettinger30c73622010-12-01 23:45:20 +0000767 def pairwise(iterable):
768 "s -> (s0,s1), (s1,s2), (s2, s3), ..."
769 a, b = tee(iterable)
770 next(b, None)
771 return zip(a, b)
Georg Brandl116aa622007-08-15 14:28:22 +0000772
Raymond Hettinger44571da2013-05-05 19:53:41 -0700773 def grouper(iterable, n, fillvalue=None):
Raymond Hettinger9ae94732012-07-08 16:04:03 -0700774 "Collect data into fixed-length chunks or blocks"
Raymond Hettinger44571da2013-05-05 19:53:41 -0700775 # grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx"
Raymond Hettinger30c73622010-12-01 23:45:20 +0000776 args = [iter(iterable)] * n
777 return zip_longest(*args, fillvalue=fillvalue)
Georg Brandl116aa622007-08-15 14:28:22 +0000778
Raymond Hettinger30c73622010-12-01 23:45:20 +0000779 def roundrobin(*iterables):
780 "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
781 # Recipe credited to George Sakkis
Raymond Hettinger337cbba2017-11-21 00:23:34 -0800782 num_active = len(iterables)
Raymond Hettinger30c73622010-12-01 23:45:20 +0000783 nexts = cycle(iter(it).__next__ for it in iterables)
Raymond Hettinger337cbba2017-11-21 00:23:34 -0800784 while num_active:
Raymond Hettinger30c73622010-12-01 23:45:20 +0000785 try:
786 for next in nexts:
787 yield next()
788 except StopIteration:
Raymond Hettinger337cbba2017-11-21 00:23:34 -0800789 # Remove the iterator we just exhausted from the cycle.
790 num_active -= 1
791 nexts = cycle(islice(nexts, num_active))
Georg Brandl116aa622007-08-15 14:28:22 +0000792
Raymond Hettinger30c73622010-12-01 23:45:20 +0000793 def partition(pred, iterable):
794 'Use a predicate to partition entries into false entries and true entries'
795 # partition(is_odd, range(10)) --> 0 2 4 6 8 and 1 3 5 7 9
796 t1, t2 = tee(iterable)
797 return filterfalse(pred, t1), filter(pred, t2)
Raymond Hettinger5ce0aa22010-12-01 10:49:19 +0000798
Raymond Hettinger30c73622010-12-01 23:45:20 +0000799 def powerset(iterable):
800 "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
801 s = list(iterable)
802 return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettinger08d01ee2010-08-07 05:36:53 +0000803
Raymond Hettinger30c73622010-12-01 23:45:20 +0000804 def unique_everseen(iterable, key=None):
805 "List unique elements, preserving order. Remember all elements ever seen."
806 # unique_everseen('AAAABBBCCDAABBB') --> A B C D
807 # unique_everseen('ABBCcAD', str.lower) --> A B C D
808 seen = set()
809 seen_add = seen.add
810 if key is None:
811 for element in filterfalse(seen.__contains__, iterable):
812 seen_add(element)
813 yield element
814 else:
815 for element in iterable:
816 k = key(element)
817 if k not in seen:
818 seen_add(k)
819 yield element
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000820
Raymond Hettinger30c73622010-12-01 23:45:20 +0000821 def unique_justseen(iterable, key=None):
822 "List unique elements, preserving order. Remember only the element just seen."
823 # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
824 # unique_justseen('ABBCcAD', str.lower) --> A B C A D
825 return map(next, map(itemgetter(1), groupby(iterable, key)))
Raymond Hettingerad9d96b2009-01-02 21:39:07 +0000826
Raymond Hettinger30c73622010-12-01 23:45:20 +0000827 def iter_except(func, exception, first=None):
828 """ Call a function repeatedly until an exception is raised.
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000829
Raymond Hettinger30c73622010-12-01 23:45:20 +0000830 Converts a call-until-exception interface to an iterator interface.
Andrew Kuchling1d7d5802013-06-20 21:17:41 -0400831 Like builtins.iter(func, sentinel) but uses an exception instead
Raymond Hettinger30c73622010-12-01 23:45:20 +0000832 of a sentinel to end the loop.
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000833
Raymond Hettinger30c73622010-12-01 23:45:20 +0000834 Examples:
835 iter_except(functools.partial(heappop, h), IndexError) # priority queue iterator
836 iter_except(d.popitem, KeyError) # non-blocking dict iterator
837 iter_except(d.popleft, IndexError) # non-blocking deque iterator
838 iter_except(q.get_nowait, Queue.Empty) # loop over a producer Queue
839 iter_except(s.pop, KeyError) # non-blocking set iterator
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000840
Raymond Hettinger30c73622010-12-01 23:45:20 +0000841 """
842 try:
843 if first is not None:
844 yield first() # For database APIs needing an initial cast to db.first()
Raymond Hettingera503f702016-03-13 00:12:31 -0800845 while True:
Raymond Hettinger30c73622010-12-01 23:45:20 +0000846 yield func()
847 except exception:
848 pass
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000849
Raymond Hettinger31b26f62014-04-02 03:16:42 -0700850 def first_true(iterable, default=False, pred=None):
851 """Returns the first true value in the iterable.
852
853 If no true value is found, returns *default*
854
855 If *pred* is not None, returns the first item
856 for which pred(item) is true.
857
858 """
859 # first_true([a,b,c], x) --> a or b or c or x
860 # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
861 return next(filter(pred, iterable), default)
862
Raymond Hettinger30c73622010-12-01 23:45:20 +0000863 def random_product(*args, repeat=1):
864 "Random selection from itertools.product(*args, **kwds)"
865 pools = [tuple(pool) for pool in args] * repeat
866 return tuple(random.choice(pool) for pool in pools)
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000867
Raymond Hettinger30c73622010-12-01 23:45:20 +0000868 def random_permutation(iterable, r=None):
869 "Random selection from itertools.permutations(iterable, r)"
870 pool = tuple(iterable)
871 r = len(pool) if r is None else r
872 return tuple(random.sample(pool, r))
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000873
Raymond Hettinger30c73622010-12-01 23:45:20 +0000874 def random_combination(iterable, r):
875 "Random selection from itertools.combinations(iterable, r)"
876 pool = tuple(iterable)
877 n = len(pool)
878 indices = sorted(random.sample(range(n), r))
879 return tuple(pool[i] for i in indices)
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000880
Raymond Hettinger30c73622010-12-01 23:45:20 +0000881 def random_combination_with_replacement(iterable, r):
882 "Random selection from itertools.combinations_with_replacement(iterable, r)"
883 pool = tuple(iterable)
884 n = len(pool)
885 indices = sorted(random.randrange(n) for i in range(r))
886 return tuple(pool[i] for i in indices)
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000887
Raymond Hettingerd37258d2018-01-13 10:35:40 -0800888 def nth_combination(iterable, r, index):
889 'Equivalent to list(combinations(iterable, r))[index]'
890 pool = tuple(iterable)
891 n = len(pool)
892 if r < 0 or r > n:
893 raise ValueError
894 c = 1
895 k = min(r, n-r)
896 for i in range(1, k+1):
897 c = c * (n - k + i) // i
898 if index < 0:
899 index += c
900 if index < 0 or index >= c:
901 raise IndexError
902 result = []
903 while r:
904 c, n, r = c*r//n, n-1, r-1
905 while index >= c:
906 index -= c
907 c, n = c*(n-r)//n, n-1
908 result.append(pool[-1-n])
909 return tuple(result)
910
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000911Note, many of the above recipes can be optimized by replacing global lookups
912with local variables defined as default values. For example, the
913*dotproduct* recipe can be written as::
914
Raymond Hettinger30c73622010-12-01 23:45:20 +0000915 def dotproduct(vec1, vec2, sum=sum, map=map, mul=operator.mul):
916 return sum(map(mul, vec1, vec2))