blob: c99ba060acf5e2413e9bb8f2d4dee24afbda808e [file] [log] [blame]
Georg Brandl116aa622007-08-15 14:28:22 +00001:mod:`itertools` --- Functions creating iterators for efficient looping
2=======================================================================
3
4.. module:: itertools
5 :synopsis: Functions creating iterators for efficient looping.
6.. moduleauthor:: Raymond Hettinger <python@rcn.com>
7.. sectionauthor:: Raymond Hettinger <python@rcn.com>
8
9
Christian Heimesfe337bf2008-03-23 21:54:12 +000010.. testsetup::
11
12 from itertools import *
13
14
Raymond Hettingerf76b9202009-02-17 20:00:59 +000015This module implements a number of :term:`iterator` building blocks inspired
16by constructs from APL, Haskell, and SML. Each has been recast in a form
17suitable for Python.
Georg Brandl116aa622007-08-15 14:28:22 +000018
19The module standardizes a core set of fast, memory efficient tools that are
Raymond Hettingerf76b9202009-02-17 20:00:59 +000020useful by themselves or in combination. Together, they form an "iterator
21algebra" making it possible to construct specialized tools succinctly and
22efficiently in pure Python.
Georg Brandl116aa622007-08-15 14:28:22 +000023
24For instance, SML provides a tabulation tool: ``tabulate(f)`` which produces a
Ezio Melottib6605992010-01-21 20:57:24 +000025sequence ``f(0), f(1), ...``. The same effect can be achieved in Python
Raymond Hettingera6c60372008-03-13 01:26:19 +000026by combining :func:`map` and :func:`count` to form ``map(f, count())``.
Georg Brandl116aa622007-08-15 14:28:22 +000027
Raymond Hettinger2c109ab2009-03-12 00:29:44 +000028These tools and their built-in counterparts also work well with the high-speed
29functions in the :mod:`operator` module. For example, the multiplication
30operator can be mapped across two vectors to form an efficient dot-product:
31``sum(map(operator.mul, vector1, vector2))``.
Georg Brandl116aa622007-08-15 14:28:22 +000032
Georg Brandl116aa622007-08-15 14:28:22 +000033
Raymond Hettingerf76b9202009-02-17 20:00:59 +000034**Infinite Iterators:**
Georg Brandl116aa622007-08-15 14:28:22 +000035
Raymond Hettinger5bfd8ce2009-04-10 19:02:36 +000036================== ================= ================================================= =========================================
37Iterator Arguments Results Example
38================== ================= ================================================= =========================================
39:func:`count` start, [step] start, start+step, start+2*step, ... ``count(10) --> 10 11 12 13 14 ...``
40:func:`cycle` p p0, p1, ... plast, p0, p1, ... ``cycle('ABCD') --> A B C D A B C D ...``
41:func:`repeat` elem [,n] elem, elem, elem, ... endlessly or up to n times ``repeat(10, 3) --> 10 10 10``
42================== ================= ================================================= =========================================
Georg Brandl116aa622007-08-15 14:28:22 +000043
Raymond Hettingerf76b9202009-02-17 20:00:59 +000044**Iterators terminating on the shortest input sequence:**
45
Raymond Hettinger5bfd8ce2009-04-10 19:02:36 +000046==================== ============================ ================================================= =============================================================
47Iterator Arguments Results Example
48==================== ============================ ================================================= =============================================================
49:func:`chain` p, q, ... p0, p1, ... plast, q0, q1, ... ``chain('ABC', 'DEF') --> A B C D E F``
50: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``
51: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``
Georg Brandlc2da5ce2009-08-13 09:16:39 +000052: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 Hettinger5bfd8ce2009-04-10 19:02:36 +000053:func:`groupby` iterable[, keyfunc] sub-iterators grouped by value of keyfunc(v)
54:func:`islice` seq, [start,] stop [, step] elements from seq[start:stop:step] ``islice('ABCDEFG', 2, None) --> C D E F G``
55:func:`starmap` func, seq func(\*seq[0]), func(\*seq[1]), ... ``starmap(pow, [(2,5), (3,2), (10,3)]) --> 32 9 1000``
56:func:`takewhile` pred, seq seq[0], seq[1], until pred fails ``takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4``
57:func:`tee` it, n it1, it2 , ... itn splits one iterator into n
Georg Brandlc2da5ce2009-08-13 09:16:39 +000058:func:`zip_longest` p, q, ... (p[0], q[0]), (p[1], q[1]), ... ``zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-``
Raymond Hettinger5bfd8ce2009-04-10 19:02:36 +000059==================== ============================ ================================================= =============================================================
Raymond Hettingerf76b9202009-02-17 20:00:59 +000060
61**Combinatoric generators:**
62
Raymond Hettinger7f587cd2009-04-10 19:43:50 +000063============================================== ==================== =============================================================
64Iterator Arguments Results
65============================================== ==================== =============================================================
66:func:`product` p, q, ... [repeat=1] cartesian product, equivalent to a nested for-loop
67:func:`permutations` p[, r] r-length tuples, all possible orderings, no repeated elements
Raymond Hettinger36c3c022009-11-19 01:20:07 +000068:func:`combinations` p, r r-length tuples, in sorted order, no repeated elements
69:func:`combinations_with_replacement` p, r r-length tuples, in sorted order, with repeated elements
Raymond Hettinger7f587cd2009-04-10 19:43:50 +000070|
71``product('ABCD', repeat=2)`` ``AA AB AC AD BA BB BC BD CA CB CC CD DA DB DC DD``
72``permutations('ABCD', 2)`` ``AB AC AD BA BC BD CA CB CD DA DB DC``
73``combinations('ABCD', 2)`` ``AB AC AD BC BD CD``
74``combinations_with_replacement('ABCD', 2)`` ``AA AB AC AD BB BC BD CC CD DD``
75============================================== ==================== =============================================================
Georg Brandl116aa622007-08-15 14:28:22 +000076
77
78.. _itertools-functions:
79
80Itertool functions
81------------------
82
83The following module functions all construct and return iterators. Some provide
84streams of infinite length, so they should only be accessed by functions or
85loops that truncate the stream.
86
87
88.. function:: chain(*iterables)
89
90 Make an iterator that returns elements from the first iterable until it is
91 exhausted, then proceeds to the next iterable, until all of the iterables are
92 exhausted. Used for treating consecutive sequences as a single sequence.
93 Equivalent to::
94
95 def chain(*iterables):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +000096 # chain('ABC', 'DEF') --> A B C D E F
Georg Brandl116aa622007-08-15 14:28:22 +000097 for it in iterables:
98 for element in it:
99 yield element
100
101
Georg Brandl933b9742010-07-29 14:36:11 +0000102.. classmethod:: chain.from_iterable(iterable)
Christian Heimes70e7ea22008-02-28 20:02:27 +0000103
Georg Brandl48310cd2009-01-03 21:18:54 +0000104 Alternate constructor for :func:`chain`. Gets chained inputs from a
Christian Heimes70e7ea22008-02-28 20:02:27 +0000105 single iterable argument that is evaluated lazily. Equivalent to::
106
107 @classmethod
108 def from_iterable(iterables):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000109 # chain.from_iterable(['ABC', 'DEF']) --> A B C D E F
Christian Heimes70e7ea22008-02-28 20:02:27 +0000110 for it in iterables:
111 for element in it:
112 yield element
113
Christian Heimes78644762008-03-04 23:39:23 +0000114
Christian Heimes836baa52008-02-26 08:18:30 +0000115.. function:: combinations(iterable, r)
116
Christian Heimesdae2a892008-04-19 00:55:37 +0000117 Return *r* length subsequences of elements from the input *iterable*.
Christian Heimes836baa52008-02-26 08:18:30 +0000118
Georg Brandl48310cd2009-01-03 21:18:54 +0000119 Combinations are emitted in lexicographic sort order. So, if the
Christian Heimes836baa52008-02-26 08:18:30 +0000120 input *iterable* is sorted, the combination tuples will be produced
Georg Brandl48310cd2009-01-03 21:18:54 +0000121 in sorted order.
Christian Heimes836baa52008-02-26 08:18:30 +0000122
123 Elements are treated as unique based on their position, not on their
124 value. So if the input elements are unique, there will be no repeat
Christian Heimes70e7ea22008-02-28 20:02:27 +0000125 values in each combination.
Christian Heimes836baa52008-02-26 08:18:30 +0000126
Christian Heimes836baa52008-02-26 08:18:30 +0000127 Equivalent to::
128
129 def combinations(iterable, r):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000130 # combinations('ABCD', 2) --> AB AC AD BC BD CD
131 # combinations(range(4), 3) --> 012 013 023 123
Christian Heimes836baa52008-02-26 08:18:30 +0000132 pool = tuple(iterable)
Christian Heimes380f7f22008-02-28 11:19:05 +0000133 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000134 if r > n:
135 return
136 indices = list(range(r))
Christian Heimesb558a2e2008-03-02 22:46:37 +0000137 yield tuple(pool[i] for i in indices)
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000138 while True:
Christian Heimes380f7f22008-02-28 11:19:05 +0000139 for i in reversed(range(r)):
Christian Heimesb558a2e2008-03-02 22:46:37 +0000140 if indices[i] != i + n - r:
Christian Heimes836baa52008-02-26 08:18:30 +0000141 break
Christian Heimes380f7f22008-02-28 11:19:05 +0000142 else:
143 return
Christian Heimesb558a2e2008-03-02 22:46:37 +0000144 indices[i] += 1
Christian Heimes380f7f22008-02-28 11:19:05 +0000145 for j in range(i+1, r):
Christian Heimesb558a2e2008-03-02 22:46:37 +0000146 indices[j] = indices[j-1] + 1
147 yield tuple(pool[i] for i in indices)
Christian Heimes836baa52008-02-26 08:18:30 +0000148
Christian Heimes78644762008-03-04 23:39:23 +0000149 The code for :func:`combinations` can be also expressed as a subsequence
150 of :func:`permutations` after filtering entries where the elements are not
151 in sorted order (according to their position in the input pool)::
152
153 def combinations(iterable, r):
154 pool = tuple(iterable)
155 n = len(pool)
156 for indices in permutations(range(n), r):
157 if sorted(indices) == list(indices):
158 yield tuple(pool[i] for i in indices)
159
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000160 The number of items returned is ``n! / r! / (n-r)!`` when ``0 <= r <= n``
161 or zero when ``r > n``.
Christian Heimes836baa52008-02-26 08:18:30 +0000162
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000163.. function:: combinations_with_replacement(iterable, r)
164
165 Return *r* length subsequences of elements from the input *iterable*
166 allowing individual elements to be repeated more than once.
167
168 Combinations are emitted in lexicographic sort order. So, if the
169 input *iterable* is sorted, the combination tuples will be produced
170 in sorted order.
171
172 Elements are treated as unique based on their position, not on their
173 value. So if the input elements are unique, the generated combinations
174 will also be unique.
175
176 Equivalent to::
177
178 def combinations_with_replacement(iterable, r):
179 # combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC
180 pool = tuple(iterable)
181 n = len(pool)
182 if not n and r:
183 return
184 indices = [0] * r
185 yield tuple(pool[i] for i in indices)
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000186 while True:
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000187 for i in reversed(range(r)):
188 if indices[i] != n - 1:
189 break
190 else:
191 return
192 indices[i:] = [indices[i] + 1] * (r - i)
193 yield tuple(pool[i] for i in indices)
194
195 The code for :func:`combinations_with_replacement` can be also expressed as
196 a subsequence of :func:`product` after filtering entries where the elements
197 are not in sorted order (according to their position in the input pool)::
198
199 def combinations_with_replacement(iterable, r):
200 pool = tuple(iterable)
201 n = len(pool)
202 for indices in product(range(n), repeat=r):
203 if sorted(indices) == list(indices):
204 yield tuple(pool[i] for i in indices)
205
206 The number of items returned is ``(n+r-1)! / r! / (n-1)!`` when ``n > 0``.
207
Raymond Hettinger30729212009-02-12 06:28:27 +0000208 .. versionadded:: 3.1
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000209
Georg Brandl67b21b72010-08-17 15:07:14 +0000210
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000211.. function:: compress(data, selectors)
212
213 Make an iterator that filters elements from *data* returning only those that
214 have a corresponding element in *selectors* that evaluates to ``True``.
Benjamin Petersond23f8222009-04-05 19:13:16 +0000215 Stops when either the *data* or *selectors* iterables has been exhausted.
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000216 Equivalent to::
217
218 def compress(data, selectors):
219 # compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F
220 return (d for d, s in zip(data, selectors) if s)
221
Raymond Hettinger30729212009-02-12 06:28:27 +0000222 .. versionadded:: 3.1
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000223
224
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000225.. function:: count(start=0, step=1)
Georg Brandl116aa622007-08-15 14:28:22 +0000226
Raymond Hettinger30729212009-02-12 06:28:27 +0000227 Make an iterator that returns evenly spaced values starting with *n*. Often
228 used as an argument to :func:`map` to generate consecutive data points.
229 Also, used with :func:`zip` to add sequence numbers. Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000230
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000231 def count(start=0, step=1):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000232 # count(10) --> 10 11 12 13 14 ...
Raymond Hettinger30729212009-02-12 06:28:27 +0000233 # count(2.5, 0.5) -> 3.5 3.0 4.5 ...
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000234 n = start
Georg Brandl116aa622007-08-15 14:28:22 +0000235 while True:
236 yield n
Raymond Hettinger30729212009-02-12 06:28:27 +0000237 n += step
Georg Brandl116aa622007-08-15 14:28:22 +0000238
Raymond Hettinger5bc472a2009-06-17 01:40:52 +0000239 When counting with floating point numbers, better accuracy can sometimes be
240 achieved by substituting multiplicative code such as: ``(start + step * i
241 for i in count())``.
242
Raymond Hettinger30729212009-02-12 06:28:27 +0000243 .. versionchanged:: 3.1
Georg Brandl67b21b72010-08-17 15:07:14 +0000244 Added *step* argument and allowed non-integer arguments.
Georg Brandl116aa622007-08-15 14:28:22 +0000245
246.. function:: cycle(iterable)
247
248 Make an iterator returning elements from the iterable and saving a copy of each.
249 When the iterable is exhausted, return elements from the saved copy. Repeats
250 indefinitely. Equivalent to::
251
252 def cycle(iterable):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000253 # cycle('ABCD') --> A B C D A B C D A B C D ...
Georg Brandl116aa622007-08-15 14:28:22 +0000254 saved = []
255 for element in iterable:
256 yield element
257 saved.append(element)
258 while saved:
259 for element in saved:
260 yield element
261
262 Note, this member of the toolkit may require significant auxiliary storage
263 (depending on the length of the iterable).
264
265
266.. function:: dropwhile(predicate, iterable)
267
268 Make an iterator that drops elements from the iterable as long as the predicate
269 is true; afterwards, returns every element. Note, the iterator does not produce
270 *any* output until the predicate first becomes false, so it may have a lengthy
271 start-up time. Equivalent to::
272
273 def dropwhile(predicate, iterable):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000274 # dropwhile(lambda x: x<5, [1,4,6,4,1]) --> 6 4 1
Georg Brandl116aa622007-08-15 14:28:22 +0000275 iterable = iter(iterable)
276 for x in iterable:
277 if not predicate(x):
278 yield x
279 break
280 for x in iterable:
281 yield x
282
Raymond Hettinger749761e2009-01-27 04:42:48 +0000283.. function:: filterfalse(predicate, iterable)
284
285 Make an iterator that filters elements from iterable returning only those for
286 which the predicate is ``False``. If *predicate* is ``None``, return the items
287 that are false. Equivalent to::
288
289 def filterfalse(predicate, iterable):
290 # filterfalse(lambda x: x%2, range(10)) --> 0 2 4 6 8
291 if predicate is None:
292 predicate = bool
293 for x in iterable:
294 if not predicate(x):
295 yield x
296
Georg Brandl116aa622007-08-15 14:28:22 +0000297
Georg Brandl3dd33882009-06-01 17:35:27 +0000298.. function:: groupby(iterable, key=None)
Georg Brandl116aa622007-08-15 14:28:22 +0000299
300 Make an iterator that returns consecutive keys and groups from the *iterable*.
301 The *key* is a function computing a key value for each element. If not
302 specified or is ``None``, *key* defaults to an identity function and returns
303 the element unchanged. Generally, the iterable needs to already be sorted on
304 the same key function.
305
306 The operation of :func:`groupby` is similar to the ``uniq`` filter in Unix. It
307 generates a break or new group every time the value of the key function changes
308 (which is why it is usually necessary to have sorted the data using the same key
309 function). That behavior differs from SQL's GROUP BY which aggregates common
310 elements regardless of their input order.
311
312 The returned group is itself an iterator that shares the underlying iterable
313 with :func:`groupby`. Because the source is shared, when the :func:`groupby`
314 object is advanced, the previous group is no longer visible. So, if that data
315 is needed later, it should be stored as a list::
316
317 groups = []
318 uniquekeys = []
319 data = sorted(data, key=keyfunc)
320 for k, g in groupby(data, keyfunc):
321 groups.append(list(g)) # Store group iterator as a list
322 uniquekeys.append(k)
323
324 :func:`groupby` is equivalent to::
325
326 class groupby(object):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000327 # [k for k, g in groupby('AAAABBBCCDAABBB')] --> A B C D A B
Raymond Hettingerd04fa312009-02-04 19:45:13 +0000328 # [list(g) for k, g in groupby('AAAABBBCCD')] --> AAAA BBB CC D
Georg Brandl116aa622007-08-15 14:28:22 +0000329 def __init__(self, iterable, key=None):
330 if key is None:
331 key = lambda x: x
332 self.keyfunc = key
333 self.it = iter(iterable)
Christian Heimes5b5e81c2007-12-31 16:14:33 +0000334 self.tgtkey = self.currkey = self.currvalue = object()
Georg Brandl116aa622007-08-15 14:28:22 +0000335 def __iter__(self):
336 return self
337 def __next__(self):
338 while self.currkey == self.tgtkey:
Raymond Hettinger21315ba2009-02-23 19:38:09 +0000339 self.currvalue = next(self.it) # Exit on StopIteration
Georg Brandl116aa622007-08-15 14:28:22 +0000340 self.currkey = self.keyfunc(self.currvalue)
341 self.tgtkey = self.currkey
342 return (self.currkey, self._grouper(self.tgtkey))
343 def _grouper(self, tgtkey):
344 while self.currkey == tgtkey:
345 yield self.currvalue
Raymond Hettinger21315ba2009-02-23 19:38:09 +0000346 self.currvalue = next(self.it) # Exit on StopIteration
Georg Brandl116aa622007-08-15 14:28:22 +0000347 self.currkey = self.keyfunc(self.currvalue)
348
Georg Brandl116aa622007-08-15 14:28:22 +0000349
Georg Brandl116aa622007-08-15 14:28:22 +0000350.. function:: islice(iterable, [start,] stop [, step])
351
352 Make an iterator that returns selected elements from the iterable. If *start* is
353 non-zero, then elements from the iterable are skipped until start is reached.
354 Afterward, elements are returned consecutively unless *step* is set higher than
355 one which results in items being skipped. If *stop* is ``None``, then iteration
356 continues until the iterator is exhausted, if at all; otherwise, it stops at the
357 specified position. Unlike regular slicing, :func:`islice` does not support
358 negative values for *start*, *stop*, or *step*. Can be used to extract related
359 fields from data where the internal structure has been flattened (for example, a
360 multi-line report may list a name field on every third line). Equivalent to::
361
362 def islice(iterable, *args):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000363 # islice('ABCDEFG', 2) --> A B
364 # islice('ABCDEFG', 2, 4) --> C D
365 # islice('ABCDEFG', 2, None) --> C D E F G
366 # islice('ABCDEFG', 0, None, 2) --> A C E G
Georg Brandl116aa622007-08-15 14:28:22 +0000367 s = slice(*args)
Raymond Hettinger7c8494b2009-05-14 21:48:02 +0000368 it = iter(range(s.start or 0, s.stop or sys.maxsize, s.step or 1))
Georg Brandl116aa622007-08-15 14:28:22 +0000369 nexti = next(it)
370 for i, element in enumerate(iterable):
371 if i == nexti:
372 yield element
373 nexti = next(it)
374
375 If *start* is ``None``, then iteration starts at zero. If *step* is ``None``,
376 then the step defaults to one.
377
Georg Brandl116aa622007-08-15 14:28:22 +0000378
Georg Brandl3dd33882009-06-01 17:35:27 +0000379.. function:: permutations(iterable, r=None)
Christian Heimes70e7ea22008-02-28 20:02:27 +0000380
381 Return successive *r* length permutations of elements in the *iterable*.
382
383 If *r* is not specified or is ``None``, then *r* defaults to the length
Georg Brandl48310cd2009-01-03 21:18:54 +0000384 of the *iterable* and all possible full-length permutations
Christian Heimes70e7ea22008-02-28 20:02:27 +0000385 are generated.
386
Georg Brandl48310cd2009-01-03 21:18:54 +0000387 Permutations are emitted in lexicographic sort order. So, if the
Christian Heimes70e7ea22008-02-28 20:02:27 +0000388 input *iterable* is sorted, the permutation tuples will be produced
Georg Brandl48310cd2009-01-03 21:18:54 +0000389 in sorted order.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000390
391 Elements are treated as unique based on their position, not on their
392 value. So if the input elements are unique, there will be no repeat
393 values in each permutation.
394
Christian Heimesb558a2e2008-03-02 22:46:37 +0000395 Equivalent to::
396
397 def permutations(iterable, r=None):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000398 # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
399 # permutations(range(3)) --> 012 021 102 120 201 210
Christian Heimesb558a2e2008-03-02 22:46:37 +0000400 pool = tuple(iterable)
401 n = len(pool)
402 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000403 if r > n:
404 return
405 indices = list(range(n))
Christian Heimesfe337bf2008-03-23 21:54:12 +0000406 cycles = range(n, n-r, -1)
Christian Heimesb558a2e2008-03-02 22:46:37 +0000407 yield tuple(pool[i] for i in indices[:r])
408 while n:
409 for i in reversed(range(r)):
410 cycles[i] -= 1
411 if cycles[i] == 0:
412 indices[i:] = indices[i+1:] + indices[i:i+1]
413 cycles[i] = n - i
414 else:
415 j = cycles[i]
416 indices[i], indices[-j] = indices[-j], indices[i]
417 yield tuple(pool[i] for i in indices[:r])
418 break
419 else:
420 return
Christian Heimes70e7ea22008-02-28 20:02:27 +0000421
Georg Brandl48310cd2009-01-03 21:18:54 +0000422 The code for :func:`permutations` can be also expressed as a subsequence of
Christian Heimes78644762008-03-04 23:39:23 +0000423 :func:`product`, filtered to exclude entries with repeated elements (those
424 from the same position in the input pool)::
425
426 def permutations(iterable, r=None):
427 pool = tuple(iterable)
428 n = len(pool)
429 r = n if r is None else r
430 for indices in product(range(n), repeat=r):
431 if len(set(indices)) == r:
432 yield tuple(pool[i] for i in indices)
433
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000434 The number of items returned is ``n! / (n-r)!`` when ``0 <= r <= n``
435 or zero when ``r > n``.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000436
Georg Brandl3dd33882009-06-01 17:35:27 +0000437.. function:: product(*iterables, repeat=1)
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000438
439 Cartesian product of input iterables.
440
441 Equivalent to nested for-loops in a generator expression. For example,
442 ``product(A, B)`` returns the same as ``((x,y) for x in A for y in B)``.
443
Christian Heimesdae2a892008-04-19 00:55:37 +0000444 The nested loops cycle like an odometer with the rightmost element advancing
445 on every iteration. This pattern creates a lexicographic ordering so that if
446 the input's iterables are sorted, the product tuples are emitted in sorted
447 order.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000448
Christian Heimes9e7f1d22008-02-28 12:27:11 +0000449 To compute the product of an iterable with itself, specify the number of
450 repetitions with the optional *repeat* keyword argument. For example,
451 ``product(A, repeat=4)`` means the same as ``product(A, A, A, A)``.
452
Christian Heimes78644762008-03-04 23:39:23 +0000453 This function is equivalent to the following code, except that the
454 actual implementation does not build up intermediate results in memory::
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000455
Raymond Hettingera137e1f2008-03-13 02:43:14 +0000456 def product(*args, repeat=1):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000457 # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy
458 # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000459 pools = [tuple(pool) for pool in args] * repeat
Christian Heimes78644762008-03-04 23:39:23 +0000460 result = [[]]
461 for pool in pools:
462 result = [x+[y] for x in result for y in pool]
463 for prod in result:
464 yield tuple(prod)
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000465
466
Raymond Hettingerd75ad442009-06-01 19:16:52 +0000467.. function:: repeat(object[, times])
Georg Brandl116aa622007-08-15 14:28:22 +0000468
469 Make an iterator that returns *object* over and over again. Runs indefinitely
Raymond Hettingera6c60372008-03-13 01:26:19 +0000470 unless the *times* argument is specified. Used as argument to :func:`map` for
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000471 invariant parameters to the called function. Also used with :func:`zip` to
Georg Brandl116aa622007-08-15 14:28:22 +0000472 create an invariant part of a tuple record. Equivalent to::
473
474 def repeat(object, times=None):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000475 # repeat(10, 3) --> 10 10 10
Georg Brandl116aa622007-08-15 14:28:22 +0000476 if times is None:
477 while True:
478 yield object
479 else:
480 for i in range(times):
481 yield object
482
483
484.. function:: starmap(function, iterable)
485
Christian Heimes679db4a2008-01-18 09:56:22 +0000486 Make an iterator that computes the function using arguments obtained from
Raymond Hettingera6c60372008-03-13 01:26:19 +0000487 the iterable. Used instead of :func:`map` when argument parameters are already
Georg Brandl116aa622007-08-15 14:28:22 +0000488 grouped in tuples from a single iterable (the data has been "pre-zipped"). The
Raymond Hettingera6c60372008-03-13 01:26:19 +0000489 difference between :func:`map` and :func:`starmap` parallels the distinction
Georg Brandl116aa622007-08-15 14:28:22 +0000490 between ``function(a,b)`` and ``function(*c)``. Equivalent to::
491
492 def starmap(function, iterable):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000493 # starmap(pow, [(2,5), (3,2), (10,3)]) --> 32 9 1000
Christian Heimes679db4a2008-01-18 09:56:22 +0000494 for args in iterable:
495 yield function(*args)
496
Georg Brandl116aa622007-08-15 14:28:22 +0000497
498.. function:: takewhile(predicate, iterable)
499
500 Make an iterator that returns elements from the iterable as long as the
501 predicate is true. Equivalent to::
502
503 def takewhile(predicate, iterable):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000504 # takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4
Georg Brandl116aa622007-08-15 14:28:22 +0000505 for x in iterable:
506 if predicate(x):
507 yield x
508 else:
509 break
510
511
Georg Brandl3dd33882009-06-01 17:35:27 +0000512.. function:: tee(iterable, n=2)
Georg Brandl116aa622007-08-15 14:28:22 +0000513
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000514 Return *n* independent iterators from a single iterable. Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000515
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000516 def tee(iterable, n=2):
517 it = iter(iterable)
518 deques = [collections.deque() for i in range(n)]
519 def gen(mydeque):
520 while True:
521 if not mydeque: # when the local deque is empty
522 newval = next(it) # fetch a new value and
523 for d in deques: # load it to all the deques
524 d.append(newval)
525 yield mydeque.popleft()
526 return tuple(gen(d) for d in deques)
Georg Brandl116aa622007-08-15 14:28:22 +0000527
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000528 Once :func:`tee` has made a split, the original *iterable* should not be
529 used anywhere else; otherwise, the *iterable* could get advanced without
530 the tee objects being informed.
Georg Brandl116aa622007-08-15 14:28:22 +0000531
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000532 This itertool may require significant auxiliary storage (depending on how
533 much temporary data needs to be stored). In general, if one iterator uses
534 most or all of the data before another iterator starts, it is faster to use
535 :func:`list` instead of :func:`tee`.
Georg Brandl116aa622007-08-15 14:28:22 +0000536
Georg Brandl116aa622007-08-15 14:28:22 +0000537
Georg Brandl3dd33882009-06-01 17:35:27 +0000538.. function:: zip_longest(*iterables, fillvalue=None)
Raymond Hettinger749761e2009-01-27 04:42:48 +0000539
540 Make an iterator that aggregates elements from each of the iterables. If the
541 iterables are of uneven length, missing values are filled-in with *fillvalue*.
542 Iteration continues until the longest iterable is exhausted. Equivalent to::
543
544 def zip_longest(*args, fillvalue=None):
545 # zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-
546 def sentinel(counter = ([fillvalue]*(len(args)-1)).pop):
547 yield counter() # yields the fillvalue, or raises IndexError
548 fillers = repeat(fillvalue)
549 iters = [chain(it, sentinel(), fillers) for it in args]
550 try:
551 for tup in zip(*iters):
552 yield tup
553 except IndexError:
554 pass
555
556 If one of the iterables is potentially infinite, then the :func:`zip_longest`
557 function should be wrapped with something that limits the number of calls
558 (for example :func:`islice` or :func:`takewhile`). If not specified,
559 *fillvalue* defaults to ``None``.
560
561
Georg Brandl116aa622007-08-15 14:28:22 +0000562.. _itertools-recipes:
563
564Recipes
565-------
566
567This section shows recipes for creating an extended toolset using the existing
568itertools as building blocks.
569
570The extended tools offer the same high performance as the underlying toolset.
571The superior memory performance is kept by processing elements one at a time
572rather than bringing the whole iterable into memory all at once. Code volume is
573kept small by linking the tools together in a functional style which helps
574eliminate temporary variables. High speed is retained by preferring
Georg Brandl9afde1c2007-11-01 20:32:30 +0000575"vectorized" building blocks over the use of for-loops and :term:`generator`\s
Christian Heimesfe337bf2008-03-23 21:54:12 +0000576which incur interpreter overhead.
577
578.. testcode::
Georg Brandl116aa622007-08-15 14:28:22 +0000579
Georg Brandl3dbca812008-07-23 16:10:53 +0000580 def take(n, iterable):
581 "Return first n items of the iterable as a list"
582 return list(islice(iterable, n))
Georg Brandl116aa622007-08-15 14:28:22 +0000583
Georg Brandl3dbca812008-07-23 16:10:53 +0000584 def tabulate(function, start=0):
Georg Brandl116aa622007-08-15 14:28:22 +0000585 "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +0000586 return map(function, count(start))
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000587
Raymond Hettingerfa007962009-03-09 11:55:25 +0000588 def consume(iterator, n):
589 "Advance the iterator n-steps ahead. If n is none, consume entirely."
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000590 # Use functions that consume iterators at C speed.
591 if n is None:
592 # feed the entire iterator into a zero-length deque
593 collections.deque(iterator, maxlen=0)
594 else:
Georg Brandla0fc3d32010-09-25 13:30:03 +0000595 # advance to the empty slice starting at position n
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000596 next(islice(iterator, n, n), None)
Raymond Hettingerfa007962009-03-09 11:55:25 +0000597
Raymond Hettingercdf8ba32009-02-19 04:45:07 +0000598 def nth(iterable, n, default=None):
599 "Returns the nth item or a default value"
600 return next(islice(iterable, n, None), default)
Georg Brandl116aa622007-08-15 14:28:22 +0000601
Georg Brandl3dbca812008-07-23 16:10:53 +0000602 def quantify(iterable, pred=bool):
603 "Count how many times the predicate is true"
604 return sum(map(pred, iterable))
Georg Brandl116aa622007-08-15 14:28:22 +0000605
Georg Brandl3dbca812008-07-23 16:10:53 +0000606 def padnone(iterable):
Georg Brandl116aa622007-08-15 14:28:22 +0000607 """Returns the sequence elements and then returns None indefinitely.
608
609 Useful for emulating the behavior of the built-in map() function.
610 """
Georg Brandl3dbca812008-07-23 16:10:53 +0000611 return chain(iterable, repeat(None))
Georg Brandl116aa622007-08-15 14:28:22 +0000612
Georg Brandl3dbca812008-07-23 16:10:53 +0000613 def ncycles(iterable, n):
Georg Brandl116aa622007-08-15 14:28:22 +0000614 "Returns the sequence elements n times"
Raymond Hettinger0f3ec6d2010-04-02 04:50:35 +0000615 return chain.from_iterable(repeat(tuple(iterable), n))
Georg Brandl116aa622007-08-15 14:28:22 +0000616
617 def dotproduct(vec1, vec2):
Raymond Hettingera6c60372008-03-13 01:26:19 +0000618 return sum(map(operator.mul, vec1, vec2))
Georg Brandl116aa622007-08-15 14:28:22 +0000619
620 def flatten(listOfLists):
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000621 "Flatten one level of nesting"
622 return chain.from_iterable(listOfLists)
Georg Brandl116aa622007-08-15 14:28:22 +0000623
624 def repeatfunc(func, times=None, *args):
625 """Repeat calls to func with specified arguments.
626
627 Example: repeatfunc(random.random)
628 """
629 if times is None:
630 return starmap(func, repeat(args))
Christian Heimes70e7ea22008-02-28 20:02:27 +0000631 return starmap(func, repeat(args, times))
Georg Brandl116aa622007-08-15 14:28:22 +0000632
633 def pairwise(iterable):
634 "s -> (s0,s1), (s1,s2), (s2, s3), ..."
635 a, b = tee(iterable)
Raymond Hettinger21315ba2009-02-23 19:38:09 +0000636 next(b, None)
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000637 return zip(a, b)
Georg Brandl116aa622007-08-15 14:28:22 +0000638
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000639 def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +0000640 "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000641 args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +0000642 return zip_longest(*args, fillvalue=fillvalue)
Georg Brandl116aa622007-08-15 14:28:22 +0000643
Christian Heimes7b3ce6a2008-01-31 14:31:45 +0000644 def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +0000645 "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimes70e7ea22008-02-28 20:02:27 +0000646 # Recipe credited to George Sakkis
Christian Heimes7b3ce6a2008-01-31 14:31:45 +0000647 pending = len(iterables)
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000648 nexts = cycle(iter(it).__next__ for it in iterables)
Christian Heimes7b3ce6a2008-01-31 14:31:45 +0000649 while pending:
650 try:
651 for next in nexts:
652 yield next()
653 except StopIteration:
654 pending -= 1
655 nexts = cycle(islice(nexts, pending))
Georg Brandl116aa622007-08-15 14:28:22 +0000656
Raymond Hettinger08d01ee2010-08-07 05:36:53 +0000657 def partition(pred, iterable):
658 'Use a predicate to partition entries into false entries and true entries'
659 # partition(is_odd, range(10)) --> 0 2 4 6 8 and 1 3 5 7 9
660 t1, t2 = tee(iterable)
661 return filterfalse(pred, t1), filter(pred, t2)
662
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000663 def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +0000664 "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
665 s = list(iterable)
666 return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000667
Raymond Hettingerfd88ea72009-03-23 22:42:28 +0000668 def unique_everseen(iterable, key=None):
669 "List unique elements, preserving order. Remember all elements ever seen."
670 # unique_everseen('AAAABBBCCDAABBB') --> A B C D
671 # unique_everseen('ABBCcAD', str.lower) --> A B C D
672 seen = set()
673 seen_add = seen.add
674 if key is None:
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000675 for element in filterfalse(seen.__contains__, iterable):
676 seen_add(element)
677 yield element
Raymond Hettingerfd88ea72009-03-23 22:42:28 +0000678 else:
679 for element in iterable:
680 k = key(element)
681 if k not in seen:
682 seen_add(k)
683 yield element
Raymond Hettingerad9d96b2009-01-02 21:39:07 +0000684
Raymond Hettingerfd88ea72009-03-23 22:42:28 +0000685 def unique_justseen(iterable, key=None):
686 "List unique elements, preserving order. Remember only the element just seen."
687 # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
688 # unique_justseen('ABBCcAD', str.lower) --> A B C A D
689 return map(next, map(itemgetter(1), groupby(iterable, key)))
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000690
691 def iter_except(func, exception, first=None):
692 """ Call a function repeatedly until an exception is raised.
693
694 Converts a call-until-exception interface to an iterator interface.
695 Like __builtin__.iter(func, sentinel) but uses an exception instead
696 of a sentinel to end the loop.
697
698 Examples:
699 iter_except(functools.partial(heappop, h), IndexError) # priority queue iterator
700 iter_except(d.popitem, KeyError) # non-blocking dict iterator
701 iter_except(d.popleft, IndexError) # non-blocking deque iterator
702 iter_except(q.get_nowait, Queue.Empty) # loop over a producer Queue
703 iter_except(s.pop, KeyError) # non-blocking set iterator
704
705 """
706 try:
707 if first is not None:
708 yield first() # For database APIs needing an initial cast to db.first()
709 while 1:
710 yield func()
711 except exception:
712 pass
713
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000714 def random_product(*args, repeat=1):
715 "Random selection from itertools.product(*args, **kwds)"
716 pools = [tuple(pool) for pool in args] * repeat
Raymond Hettinger0f3ec6d2010-04-02 04:50:35 +0000717 return tuple(random.choice(pool) for pool in pools)
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000718
Raymond Hettingerc075f072010-04-10 07:09:53 +0000719 def random_permutation(iterable, r=None):
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000720 "Random selection from itertools.permutations(iterable, r)"
721 pool = tuple(iterable)
722 r = len(pool) if r is None else r
Raymond Hettinger0f3ec6d2010-04-02 04:50:35 +0000723 return tuple(random.sample(pool, r))
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000724
725 def random_combination(iterable, r):
726 "Random selection from itertools.combinations(iterable, r)"
727 pool = tuple(iterable)
Raymond Hettingerc075f072010-04-10 07:09:53 +0000728 n = len(pool)
729 indices = sorted(random.sample(range(n), r))
730 return tuple(pool[i] for i in indices)
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000731
732 def random_combination_with_replacement(iterable, r):
733 "Random selection from itertools.combinations_with_replacement(iterable, r)"
734 pool = tuple(iterable)
Raymond Hettingerc075f072010-04-10 07:09:53 +0000735 n = len(pool)
736 indices = sorted(random.randrange(n) for i in range(r))
737 return tuple(pool[i] for i in indices)
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000738
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000739Note, many of the above recipes can be optimized by replacing global lookups
740with local variables defined as default values. For example, the
741*dotproduct* recipe can be written as::
742
743 def dotproduct(vec1, vec2, sum=sum, map=map, mul=operator.mul):
744 return sum(map(mul, vec1, vec2))