]> git.ipfire.org Git - thirdparty/Python/cpython.git/blame - Doc/library/itertools.rst
Misc itertool recipe improvements, mostly docstrings and comments (gh-109555)
[thirdparty/Python/cpython.git] / Doc / library / itertools.rst
CommitLineData
116aa62b
GB
1:mod:`itertools` --- Functions creating iterators for efficient looping
2=======================================================================
3
4.. module:: itertools
30c73624 5 :synopsis: Functions creating iterators for efficient looping.
fa089b9b 6
116aa62b
GB
7.. moduleauthor:: Raymond Hettinger <python@rcn.com>
8.. sectionauthor:: Raymond Hettinger <python@rcn.com>
9
fe337bfd
CH
10.. testsetup::
11
30c73624 12 from itertools import *
f4ead487
RH
13 import collections
14 import math
15 import operator
16 import random
fe337bfd 17
fa089b9b 18--------------
fe337bfd 19
f76b9209
RH
20This module implements a number of :term:`iterator` building blocks inspired
21by constructs from APL, Haskell, and SML. Each has been recast in a form
22suitable for Python.
116aa62b
GB
23
24The module standardizes a core set of fast, memory efficient tools that are
f76b9209
RH
25useful by themselves or in combination. Together, they form an "iterator
26algebra" making it possible to construct specialized tools succinctly and
27efficiently in pure Python.
116aa62b
GB
28
29For instance, SML provides a tabulation tool: ``tabulate(f)`` which produces a
b660599f 30sequence ``f(0), f(1), ...``. The same effect can be achieved in Python
a6c6037f 31by combining :func:`map` and :func:`count` to form ``map(f, count())``.
116aa62b 32
2c109ab5
RH
33These tools and their built-in counterparts also work well with the high-speed
34functions in the :mod:`operator` module. For example, the multiplication
35operator can be mapped across two vectors to form an efficient dot-product:
47b9f83a 36``sum(starmap(operator.mul, zip(vec1, vec2, strict=True)))``.
f76b9209
RH
37
38
6693d7af 39**Infinite iterators:**
f76b9209 40
5bfd8ce5
RH
41================== ================= ================================================= =========================================
42Iterator Arguments Results Example
43================== ================= ================================================= =========================================
44:func:`count` start, [step] start, start+step, start+2*step, ... ``count(10) --> 10 11 12 13 14 ...``
45:func:`cycle` p p0, p1, ... plast, p0, p1, ... ``cycle('ABCD') --> A B C D A B C D ...``
46:func:`repeat` elem [,n] elem, elem, elem, ... endlessly or up to n times ``repeat(10, 3) --> 10 10 10``
47================== ================= ================================================= =========================================
f76b9209
RH
48
49**Iterators terminating on the shortest input sequence:**
50
2989f587
BP
51============================ ============================ ================================================= =============================================================
52Iterator Arguments Results Example
53============================ ============================ ================================================= =============================================================
54:func:`accumulate` p [,func] p0, p0+p1, p0+p1+p2, ... ``accumulate([1,2,3,4,5]) --> 1 3 6 10 15``
35cc0ea7 55:func:`batched` p, n (p0, p1, ..., p_n-1), ... ``batched('ABCDEFG', n=3) --> ABC DEF G``
2989f587
BP
56:func:`chain` p, q, ... p0, p1, ... plast, q0, q1, ... ``chain('ABC', 'DEF') --> A B C D E F``
57:func:`chain.from_iterable` iterable p0, p1, ... plast, q0, q1, ... ``chain.from_iterable(['ABC', 'DEF']) --> A B C D E F``
58: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``
59: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``
60:func:`filterfalse` pred, seq elements of seq where pred(elem) is false ``filterfalse(lambda x: x%2, range(10)) --> 0 2 4 6 8``
49392c63 61:func:`groupby` iterable[, key] sub-iterators grouped by value of key(v)
2989f587 62:func:`islice` seq, [start,] stop [, step] elements from seq[start:stop:step] ``islice('ABCDEFG', 2, None) --> C D E F G``
cc061d0e 63:func:`pairwise` iterable (p[0], p[1]), (p[1], p[2]) ``pairwise('ABCDEFG') --> AB BC CD DE EF FG``
2989f587
BP
64:func:`starmap` func, seq func(\*seq[0]), func(\*seq[1]), ... ``starmap(pow, [(2,5), (3,2), (10,3)]) --> 32 9 1000``
65:func:`takewhile` pred, seq seq[0], seq[1], until pred fails ``takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4``
66:func:`tee` it, n it1, it2, ... itn splits one iterator into n
67:func:`zip_longest` p, q, ... (p[0], q[0]), (p[1], q[1]), ... ``zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-``
68============================ ============================ ================================================= =============================================================
f76b9209 69
6693d7af 70**Combinatoric iterators:**
f76b9209 71
7f587cd3
RH
72============================================== ==================== =============================================================
73Iterator Arguments Results
74============================================== ==================== =============================================================
75:func:`product` p, q, ... [repeat=1] cartesian product, equivalent to a nested for-loop
76:func:`permutations` p[, r] r-length tuples, all possible orderings, no repeated elements
36c3c027
RH
77:func:`combinations` p, r r-length tuples, in sorted order, no repeated elements
78:func:`combinations_with_replacement` p, r r-length tuples, in sorted order, with repeated elements
7f587cd3 79============================================== ==================== =============================================================
116aa62b 80
14e3c447
BW
81============================================== =============================================================
82Examples Results
83============================================== =============================================================
84``product('ABCD', repeat=2)`` ``AA AB AC AD BA BB BC BD CA CB CC CD DA DB DC DD``
85``permutations('ABCD', 2)`` ``AB AC AD BA BC BD CA CB CD DA DB DC``
86``combinations('ABCD', 2)`` ``AB AC AD BC BD CD``
87``combinations_with_replacement('ABCD', 2)`` ``AA AB AC AD BB BC BD CC CD DD``
88============================================== =============================================================
89
116aa62b
GB
90
91.. _itertools-functions:
92
93Itertool functions
94------------------
95
96The following module functions all construct and return iterators. Some provide
97streams of infinite length, so they should only be accessed by functions or
98loops that truncate the stream.
99
9718b59e 100.. function:: accumulate(iterable[, func, *, initial=None])
adb8146e 101
15b04eb4
AK
102 Make an iterator that returns accumulated sums, or accumulated
103 results of other binary functions (specified via the optional
9718b59e
LR
104 *func* argument).
105
106 If *func* is supplied, it should be a function
15b04eb4
AK
107 of two arguments. Elements of the input *iterable* may be any type
108 that can be accepted as arguments to *func*. (For example, with
109 the default operation of addition, elements may be any addable
110 type including :class:`~decimal.Decimal` or
9718b59e
LR
111 :class:`~fractions.Fraction`.)
112
113 Usually, the number of elements output matches the input iterable.
114 However, if the keyword argument *initial* is provided, the
115 accumulation leads off with the *initial* value so that the output
116 has one more element than the input iterable.
adb8146e 117
672866d0 118 Roughly equivalent to::
5d44613e 119
9718b59e 120 def accumulate(iterable, func=operator.add, *, initial=None):
adb8146e 121 'Return running totals'
d8ff4658 122 # accumulate([1,2,3,4,5]) --> 1 3 6 10 15
9718b59e 123 # accumulate([1,2,3,4,5], initial=100) --> 100 101 103 106 110 115
5d44613e 124 # accumulate([1,2,3,4,5], operator.mul) --> 1 2 6 24 120
d8ff4658 125 it = iter(iterable)
9718b59e
LR
126 total = initial
127 if initial is None:
128 try:
129 total = next(it)
130 except StopIteration:
131 return
d8ff4658
RH
132 yield total
133 for element in it:
5d44613e 134 total = func(total, element)
d8ff4658 135 yield total
adb8146e 136
295c1d4f
RH
137 There are a number of uses for the *func* argument. It can be set to
138 :func:`min` for a running minimum, :func:`max` for a running maximum, or
139 :func:`operator.mul` for a running product. Amortization tables can be
f4ead487
RH
140 built by accumulating interest and applying payments:
141
142 .. doctest::
5d44613e
RH
143
144 >>> data = [3, 4, 6, 2, 1, 9, 0, 7, 5, 8]
145 >>> list(accumulate(data, operator.mul)) # running product
146 [3, 12, 72, 144, 144, 1296, 0, 0, 0, 0]
147 >>> list(accumulate(data, max)) # running maximum
148 [3, 4, 6, 6, 6, 9, 9, 9, 9, 9]
149
278030a1
RH
150 # Amortize a 5% loan of 1000 with 10 annual payments of 90
151 >>> account_update = lambda bal, pmt: round(bal * 1.05) + pmt
152 >>> list(accumulate(repeat(-90, 10), account_update, initial=1_000))
153 [1000, 960, 918, 874, 828, 779, 728, 674, 618, 559, 497]
5d44613e 154
64801680
RH
155 See :func:`functools.reduce` for a similar function that returns only the
156 final accumulated value.
157
adb8146e 158 .. versionadded:: 3.2
116aa62b 159
5d44613e
RH
160 .. versionchanged:: 3.3
161 Added the optional *func* parameter.
162
9718b59e
LR
163 .. versionchanged:: 3.8
164 Added the optional *initial* parameter.
165
de3ece76
RH
166
167.. function:: batched(iterable, n)
168
35cc0ea7 169 Batch data from the *iterable* into tuples of length *n*. The last
de3ece76
RH
170 batch may be shorter than *n*.
171
35cc0ea7
RH
172 Loops over the input iterable and accumulates data into tuples up to
173 size *n*. The input is consumed lazily, just enough to fill a batch.
de3ece76
RH
174 The result is yielded as soon as the batch is full or when the input
175 iterable is exhausted:
176
177 .. doctest::
178
179 >>> flattened_data = ['roses', 'red', 'violets', 'blue', 'sugar', 'sweet']
180 >>> unflattened = list(batched(flattened_data, 2))
181 >>> unflattened
35cc0ea7 182 [('roses', 'red'), ('violets', 'blue'), ('sugar', 'sweet')]
de3ece76
RH
183
184 >>> for batch in batched('ABCDEFG', 3):
185 ... print(batch)
186 ...
35cc0ea7
RH
187 ('A', 'B', 'C')
188 ('D', 'E', 'F')
189 ('G',)
de3ece76
RH
190
191 Roughly equivalent to::
192
193 def batched(iterable, n):
194 # batched('ABCDEFG', 3) --> ABC DEF G
195 if n < 1:
196 raise ValueError('n must be at least one')
197 it = iter(iterable)
4bb1dd3c 198 while batch := tuple(islice(it, n)):
de3ece76
RH
199 yield batch
200
a53f6373 201 .. versionadded:: 3.12
de3ece76
RH
202
203
116aa62b
GB
204.. function:: chain(*iterables)
205
30c73624
RH
206 Make an iterator that returns elements from the first iterable until it is
207 exhausted, then proceeds to the next iterable, until all of the iterables are
208 exhausted. Used for treating consecutive sequences as a single sequence.
672866d0 209 Roughly equivalent to::
116aa62b 210
30c73624
RH
211 def chain(*iterables):
212 # chain('ABC', 'DEF') --> A B C D E F
213 for it in iterables:
214 for element in it:
215 yield element
116aa62b
GB
216
217
933b974a 218.. classmethod:: chain.from_iterable(iterable)
70e7ea23 219
30c73624 220 Alternate constructor for :func:`chain`. Gets chained inputs from a
1e21ebcc 221 single iterable argument that is evaluated lazily. Roughly equivalent to::
70e7ea23 222
30c73624
RH
223 def from_iterable(iterables):
224 # chain.from_iterable(['ABC', 'DEF']) --> A B C D E F
225 for it in iterables:
226 for element in it:
227 yield element
70e7ea23 228
7864476a 229
836baa53
CH
230.. function:: combinations(iterable, r)
231
30c73624 232 Return *r* length subsequences of elements from the input *iterable*.
836baa53 233
5e0ed8ab
RW
234 The combination tuples are emitted in lexicographic ordering according to
235 the order of the input *iterable*. So, if the input *iterable* is sorted,
f4ead487 236 the output tuples will be produced in sorted order.
836baa53 237
30c73624 238 Elements are treated as unique based on their position, not on their
f4ead487 239 value. So if the input elements are unique, there will be no repeated
30c73624 240 values in each combination.
836baa53 241
672866d0 242 Roughly equivalent to::
836baa53
CH
243
244 def combinations(iterable, r):
dd1150e3
RH
245 # combinations('ABCD', 2) --> AB AC AD BC BD CD
246 # combinations(range(4), 3) --> 012 013 023 123
836baa53 247 pool = tuple(iterable)
380f7f22 248 n = len(pool)
5bad41ee
RH
249 if r > n:
250 return
251 indices = list(range(r))
b558a2e1 252 yield tuple(pool[i] for i in indices)
cf984cee 253 while True:
380f7f22 254 for i in reversed(range(r)):
b558a2e1 255 if indices[i] != i + n - r:
836baa53 256 break
380f7f22
CH
257 else:
258 return
b558a2e1 259 indices[i] += 1
380f7f22 260 for j in range(i+1, r):
b558a2e1
CH
261 indices[j] = indices[j-1] + 1
262 yield tuple(pool[i] for i in indices)
836baa53 263
30c73624
RH
264 The code for :func:`combinations` can be also expressed as a subsequence
265 of :func:`permutations` after filtering entries where the elements are not
266 in sorted order (according to their position in the input pool)::
7864476a
CH
267
268 def combinations(iterable, r):
269 pool = tuple(iterable)
270 n = len(pool)
271 for indices in permutations(range(n), r):
272 if sorted(indices) == list(indices):
273 yield tuple(pool[i] for i in indices)
274
30c73624
RH
275 The number of items returned is ``n! / r! / (n-r)!`` when ``0 <= r <= n``
276 or zero when ``r > n``.
836baa53 277
d07d939c
RH
278.. function:: combinations_with_replacement(iterable, r)
279
30c73624
RH
280 Return *r* length subsequences of elements from the input *iterable*
281 allowing individual elements to be repeated more than once.
d07d939c 282
5e0ed8ab
RW
283 The combination tuples are emitted in lexicographic ordering according to
284 the order of the input *iterable*. So, if the input *iterable* is sorted,
f4ead487 285 the output tuples will be produced in sorted order.
d07d939c 286
30c73624
RH
287 Elements are treated as unique based on their position, not on their
288 value. So if the input elements are unique, the generated combinations
289 will also be unique.
d07d939c 290
672866d0 291 Roughly equivalent to::
d07d939c
RH
292
293 def combinations_with_replacement(iterable, r):
294 # combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC
295 pool = tuple(iterable)
296 n = len(pool)
297 if not n and r:
298 return
299 indices = [0] * r
300 yield tuple(pool[i] for i in indices)
cf984cee 301 while True:
d07d939c
RH
302 for i in reversed(range(r)):
303 if indices[i] != n - 1:
304 break
305 else:
306 return
307 indices[i:] = [indices[i] + 1] * (r - i)
308 yield tuple(pool[i] for i in indices)
309
30c73624
RH
310 The code for :func:`combinations_with_replacement` can be also expressed as
311 a subsequence of :func:`product` after filtering entries where the elements
312 are not in sorted order (according to their position in the input pool)::
d07d939c
RH
313
314 def combinations_with_replacement(iterable, r):
315 pool = tuple(iterable)
316 n = len(pool)
317 for indices in product(range(n), repeat=r):
318 if sorted(indices) == list(indices):
319 yield tuple(pool[i] for i in indices)
320
30c73624 321 The number of items returned is ``(n+r-1)! / r! / (n-1)!`` when ``n > 0``.
d07d939c 322
30c73624 323 .. versionadded:: 3.1
d07d939c 324
67b21b75 325
6b3b0fc4
RH
326.. function:: compress(data, selectors)
327
30c73624
RH
328 Make an iterator that filters elements from *data* returning only those that
329 have a corresponding element in *selectors* that evaluates to ``True``.
330 Stops when either the *data* or *selectors* iterables has been exhausted.
672866d0 331 Roughly equivalent to::
6b3b0fc4 332
30c73624
RH
333 def compress(data, selectors):
334 # compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F
335 return (d for d, s in zip(data, selectors) if s)
6b3b0fc4 336
30c73624 337 .. versionadded:: 3.1
6b3b0fc4
RH
338
339
9e8dbbcd 340.. function:: count(start=0, step=1)
116aa62b 341
edb4260f 342 Make an iterator that returns evenly spaced values starting with number *start*. Often
30c73624 343 used as an argument to :func:`map` to generate consecutive data points.
672866d0 344 Also, used with :func:`zip` to add sequence numbers. Roughly equivalent to::
116aa62b 345
30c73624
RH
346 def count(start=0, step=1):
347 # count(10) --> 10 11 12 13 14 ...
8dc9b3fb 348 # count(2.5, 0.5) --> 2.5 3.0 3.5 ...
30c73624
RH
349 n = start
350 while True:
351 yield n
352 n += step
116aa62b 353
30c73624
RH
354 When counting with floating point numbers, better accuracy can sometimes be
355 achieved by substituting multiplicative code such as: ``(start + step * i
356 for i in count())``.
5bc472a9 357
30c73624
RH
358 .. versionchanged:: 3.1
359 Added *step* argument and allowed non-integer arguments.
116aa62b
GB
360
361.. function:: cycle(iterable)
362
30c73624
RH
363 Make an iterator returning elements from the iterable and saving a copy of each.
364 When the iterable is exhausted, return elements from the saved copy. Repeats
672866d0 365 indefinitely. Roughly equivalent to::
30c73624
RH
366
367 def cycle(iterable):
368 # cycle('ABCD') --> A B C D A B C D A B C D ...
369 saved = []
370 for element in iterable:
371 yield element
372 saved.append(element)
373 while saved:
374 for element in saved:
116aa62b
GB
375 yield element
376
30c73624
RH
377 Note, this member of the toolkit may require significant auxiliary storage
378 (depending on the length of the iterable).
116aa62b
GB
379
380
381.. function:: dropwhile(predicate, iterable)
382
30c73624
RH
383 Make an iterator that drops elements from the iterable as long as the predicate
384 is true; afterwards, returns every element. Note, the iterator does not produce
385 *any* output until the predicate first becomes false, so it may have a lengthy
672866d0 386 start-up time. Roughly equivalent to::
30c73624
RH
387
388 def dropwhile(predicate, iterable):
389 # dropwhile(lambda x: x<5, [1,4,6,4,1]) --> 6 4 1
390 iterable = iter(iterable)
391 for x in iterable:
392 if not predicate(x):
393 yield x
394 break
395 for x in iterable:
396 yield x
116aa62b 397
749761e1
RH
398.. function:: filterfalse(predicate, iterable)
399
30c73624 400 Make an iterator that filters elements from iterable returning only those for
81bf10e4 401 which the predicate is false. If *predicate* is ``None``, return the items
672866d0 402 that are false. Roughly equivalent to::
749761e1 403
30c73624
RH
404 def filterfalse(predicate, iterable):
405 # filterfalse(lambda x: x%2, range(10)) --> 0 2 4 6 8
406 if predicate is None:
407 predicate = bool
408 for x in iterable:
409 if not predicate(x):
410 yield x
749761e1 411
116aa62b 412
3dd33882 413.. function:: groupby(iterable, key=None)
116aa62b 414
30c73624
RH
415 Make an iterator that returns consecutive keys and groups from the *iterable*.
416 The *key* is a function computing a key value for each element. If not
417 specified or is ``None``, *key* defaults to an identity function and returns
418 the element unchanged. Generally, the iterable needs to already be sorted on
419 the same key function.
420
421 The operation of :func:`groupby` is similar to the ``uniq`` filter in Unix. It
422 generates a break or new group every time the value of the key function changes
423 (which is why it is usually necessary to have sorted the data using the same key
424 function). That behavior differs from SQL's GROUP BY which aggregates common
425 elements regardless of their input order.
426
427 The returned group is itself an iterator that shares the underlying iterable
428 with :func:`groupby`. Because the source is shared, when the :func:`groupby`
429 object is advanced, the previous group is no longer visible. So, if that data
430 is needed later, it should be stored as a list::
431
432 groups = []
433 uniquekeys = []
434 data = sorted(data, key=keyfunc)
435 for k, g in groupby(data, keyfunc):
436 groups.append(list(g)) # Store group iterator as a list
437 uniquekeys.append(k)
438
672866d0 439 :func:`groupby` is roughly equivalent to::
30c73624
RH
440
441 class groupby:
442 # [k for k, g in groupby('AAAABBBCCDAABBB')] --> A B C D A B
443 # [list(g) for k, g in groupby('AAAABBBCCD')] --> AAAA BBB CC D
f4ead487 444
30c73624
RH
445 def __init__(self, iterable, key=None):
446 if key is None:
447 key = lambda x: x
448 self.keyfunc = key
449 self.it = iter(iterable)
450 self.tgtkey = self.currkey = self.currvalue = object()
f4ead487 451
30c73624
RH
452 def __iter__(self):
453 return self
f4ead487 454
30c73624 455 def __next__(self):
c247caf3 456 self.id = object()
30c73624
RH
457 while self.currkey == self.tgtkey:
458 self.currvalue = next(self.it) # Exit on StopIteration
459 self.currkey = self.keyfunc(self.currvalue)
460 self.tgtkey = self.currkey
c247caf3 461 return (self.currkey, self._grouper(self.tgtkey, self.id))
f4ead487 462
c247caf3
SS
463 def _grouper(self, tgtkey, id):
464 while self.id is id and self.currkey == tgtkey:
30c73624 465 yield self.currvalue
828d932a
RH
466 try:
467 self.currvalue = next(self.it)
468 except StopIteration:
469 return
30c73624 470 self.currkey = self.keyfunc(self.currvalue)
116aa62b 471
116aa62b 472
e0add764
EM
473.. function:: islice(iterable, stop)
474 islice(iterable, start, stop[, step])
116aa62b 475
30c73624
RH
476 Make an iterator that returns selected elements from the iterable. If *start* is
477 non-zero, then elements from the iterable are skipped until start is reached.
478 Afterward, elements are returned consecutively unless *step* is set higher than
479 one which results in items being skipped. If *stop* is ``None``, then iteration
480 continues until the iterator is exhausted, if at all; otherwise, it stops at the
f4ead487
RH
481 specified position.
482
483 If *start* is ``None``, then iteration starts at zero. If *step* is ``None``,
484 then the step defaults to one.
485
486 Unlike regular slicing, :func:`islice` does not support negative values for
487 *start*, *stop*, or *step*. Can be used to extract related fields from
488 data where the internal structure has been flattened (for example, a
489 multi-line report may list a name field on every third line).
490
491 Roughly equivalent to::
30c73624
RH
492
493 def islice(iterable, *args):
494 # islice('ABCDEFG', 2) --> A B
495 # islice('ABCDEFG', 2, 4) --> C D
496 # islice('ABCDEFG', 2, None) --> C D E F G
497 # islice('ABCDEFG', 0, None, 2) --> A C E G
498 s = slice(*args)
da1734c5
CS
499 start, stop, step = s.start or 0, s.stop or sys.maxsize, s.step or 1
500 it = iter(range(start, stop, step))
828d932a
RH
501 try:
502 nexti = next(it)
503 except StopIteration:
da1734c5
CS
504 # Consume *iterable* up to the *start* position.
505 for i, element in zip(range(start), iterable):
506 pass
828d932a 507 return
da1734c5
CS
508 try:
509 for i, element in enumerate(iterable):
510 if i == nexti:
511 yield element
512 nexti = next(it)
513 except StopIteration:
514 # Consume to *stop*.
515 for i, element in zip(range(i + 1, stop), iterable):
516 pass
30c73624 517
116aa62b 518
cc061d0e
RH
519.. function:: pairwise(iterable)
520
521 Return successive overlapping pairs taken from the input *iterable*.
522
523 The number of 2-tuples in the output iterator will be one fewer than the
524 number of inputs. It will be empty if the input iterable has fewer than
525 two values.
526
527 Roughly equivalent to::
528
529 def pairwise(iterable):
530 # pairwise('ABCDEFG') --> AB BC CD DE EF FG
531 a, b = tee(iterable)
532 next(b, None)
533 return zip(a, b)
534
750368cb
RH
535 .. versionadded:: 3.10
536
116aa62b 537
3dd33882 538.. function:: permutations(iterable, r=None)
70e7ea23 539
30c73624 540 Return successive *r* length permutations of elements in the *iterable*.
70e7ea23 541
30c73624
RH
542 If *r* is not specified or is ``None``, then *r* defaults to the length
543 of the *iterable* and all possible full-length permutations
544 are generated.
70e7ea23 545
f4ead487 546 The permutation tuples are emitted in lexicographic order according to
5e0ed8ab 547 the order of the input *iterable*. So, if the input *iterable* is sorted,
f4ead487 548 the output tuples will be produced in sorted order.
70e7ea23 549
30c73624 550 Elements are treated as unique based on their position, not on their
f4ead487
RH
551 value. So if the input elements are unique, there will be no repeated
552 values within a permutation.
70e7ea23 553
672866d0 554 Roughly equivalent to::
b558a2e1
CH
555
556 def permutations(iterable, r=None):
dd1150e3
RH
557 # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
558 # permutations(range(3)) --> 012 021 102 120 201 210
b558a2e1
CH
559 pool = tuple(iterable)
560 n = len(pool)
561 r = n if r is None else r
5bad41ee
RH
562 if r > n:
563 return
564 indices = list(range(n))
73866626 565 cycles = list(range(n, n-r, -1))
b558a2e1
CH
566 yield tuple(pool[i] for i in indices[:r])
567 while n:
568 for i in reversed(range(r)):
569 cycles[i] -= 1
570 if cycles[i] == 0:
571 indices[i:] = indices[i+1:] + indices[i:i+1]
572 cycles[i] = n - i
573 else:
574 j = cycles[i]
575 indices[i], indices[-j] = indices[-j], indices[i]
576 yield tuple(pool[i] for i in indices[:r])
577 break
578 else:
579 return
70e7ea23 580
30c73624
RH
581 The code for :func:`permutations` can be also expressed as a subsequence of
582 :func:`product`, filtered to exclude entries with repeated elements (those
583 from the same position in the input pool)::
7864476a
CH
584
585 def permutations(iterable, r=None):
586 pool = tuple(iterable)
587 n = len(pool)
588 r = n if r is None else r
589 for indices in product(range(n), repeat=r):
590 if len(set(indices)) == r:
591 yield tuple(pool[i] for i in indices)
592
30c73624
RH
593 The number of items returned is ``n! / (n-r)!`` when ``0 <= r <= n``
594 or zero when ``r > n``.
70e7ea23 595
3dd33882 596.. function:: product(*iterables, repeat=1)
90c3d9b9 597
30c73624 598 Cartesian product of input iterables.
90c3d9b9 599
672866d0 600 Roughly equivalent to nested for-loops in a generator expression. For example,
30c73624 601 ``product(A, B)`` returns the same as ``((x,y) for x in A for y in B)``.
90c3d9b9 602
30c73624
RH
603 The nested loops cycle like an odometer with the rightmost element advancing
604 on every iteration. This pattern creates a lexicographic ordering so that if
605 the input's iterables are sorted, the product tuples are emitted in sorted
606 order.
90c3d9b9 607
30c73624
RH
608 To compute the product of an iterable with itself, specify the number of
609 repetitions with the optional *repeat* keyword argument. For example,
610 ``product(A, repeat=4)`` means the same as ``product(A, A, A, A)``.
9e7f1d2e 611
672866d0 612 This function is roughly equivalent to the following code, except that the
30c73624 613 actual implementation does not build up intermediate results in memory::
90c3d9b9 614
30c73624
RH
615 def product(*args, repeat=1):
616 # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy
617 # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111
618 pools = [tuple(pool) for pool in args] * repeat
619 result = [[]]
620 for pool in pools:
621 result = [x+[y] for x in result for y in pool]
622 for prod in result:
623 yield tuple(prod)
90c3d9b9 624
cfc6ce4d
RN
625 Before :func:`product` runs, it completely consumes the input iterables,
626 keeping pools of values in memory to generate the products. Accordingly,
074ad512 627 it is only useful with finite inputs.
90c3d9b9 628
d75ad44a 629.. function:: repeat(object[, times])
116aa62b 630
30c73624 631 Make an iterator that returns *object* over and over again. Runs indefinitely
f4ead487 632 unless the *times* argument is specified.
672866d0
RH
633
634 Roughly equivalent to::
116aa62b 635
30c73624
RH
636 def repeat(object, times=None):
637 # repeat(10, 3) --> 10 10 10
638 if times is None:
639 while True:
640 yield object
641 else:
642 for i in range(times):
643 yield object
116aa62b 644
fc3ba6b8 645 A common use for *repeat* is to supply a stream of constant values to *map*
f4ead487
RH
646 or *zip*:
647
648 .. doctest::
fc3ba6b8
RH
649
650 >>> list(map(pow, range(10), repeat(2)))
651 [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
116aa62b
GB
652
653.. function:: starmap(function, iterable)
654
30c73624
RH
655 Make an iterator that computes the function using arguments obtained from
656 the iterable. Used instead of :func:`map` when argument parameters are already
f4ead487
RH
657 grouped in tuples from a single iterable (when the data has been
658 "pre-zipped").
659
660 The difference between :func:`map` and :func:`starmap` parallels the
661 distinction between ``function(a,b)`` and ``function(*c)``. Roughly
662 equivalent to::
116aa62b 663
30c73624
RH
664 def starmap(function, iterable):
665 # starmap(pow, [(2,5), (3,2), (10,3)]) --> 32 9 1000
666 for args in iterable:
667 yield function(*args)
679db4aa 668
116aa62b
GB
669
670.. function:: takewhile(predicate, iterable)
671
30c73624 672 Make an iterator that returns elements from the iterable as long as the
672866d0 673 predicate is true. Roughly equivalent to::
116aa62b 674
30c73624
RH
675 def takewhile(predicate, iterable):
676 # takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4
677 for x in iterable:
678 if predicate(x):
679 yield x
680 else:
681 break
116aa62b
GB
682
683
3dd33882 684.. function:: tee(iterable, n=2)
116aa62b 685
ab425aa9
RH
686 Return *n* independent iterators from a single iterable.
687
688 The following Python code helps explain what *tee* does (although the actual
4ecfa455 689 implementation is more complex and uses only a single underlying
f4ead487 690 :abbr:`FIFO (first-in, first-out)` queue)::
cf984cee
RH
691
692 def tee(iterable, n=2):
693 it = iter(iterable)
694 deques = [collections.deque() for i in range(n)]
695 def gen(mydeque):
696 while True:
697 if not mydeque: # when the local deque is empty
828d932a
RH
698 try:
699 newval = next(it) # fetch a new value and
700 except StopIteration:
701 return
cf984cee
RH
702 for d in deques: # load it to all the deques
703 d.append(newval)
704 yield mydeque.popleft()
705 return tuple(gen(d) for d in deques)
706
f4ead487 707 Once a :func:`tee` has been created, the original *iterable* should not be
30c73624 708 used anywhere else; otherwise, the *iterable* could get advanced without
918b468b 709 the tee objects being informed.
cf984cee 710
526a0146 711 ``tee`` iterators are not threadsafe. A :exc:`RuntimeError` may be
757c383d 712 raised when simultaneously using iterators returned by the same :func:`tee`
526a0146
SS
713 call, even if the original *iterable* is threadsafe.
714
30c73624
RH
715 This itertool may require significant auxiliary storage (depending on how
716 much temporary data needs to be stored). In general, if one iterator uses
717 most or all of the data before another iterator starts, it is faster to use
718 :func:`list` instead of :func:`tee`.
116aa62b 719
116aa62b 720
3dd33882 721.. function:: zip_longest(*iterables, fillvalue=None)
749761e1 722
30c73624
RH
723 Make an iterator that aggregates elements from each of the iterables. If the
724 iterables are of uneven length, missing values are filled-in with *fillvalue*.
672866d0 725 Iteration continues until the longest iterable is exhausted. Roughly equivalent to::
749761e1 726
3147b042 727 def zip_longest(*args, fillvalue=None):
30c73624 728 # zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-
3147b042
RH
729 iterators = [iter(it) for it in args]
730 num_active = len(iterators)
731 if not num_active:
732 return
733 while True:
734 values = []
735 for i, it in enumerate(iterators):
736 try:
737 value = next(it)
738 except StopIteration:
739 num_active -= 1
740 if not num_active:
741 return
742 iterators[i] = repeat(fillvalue)
743 value = fillvalue
744 values.append(value)
745 yield tuple(values)
749761e1 746
30c73624
RH
747 If one of the iterables is potentially infinite, then the :func:`zip_longest`
748 function should be wrapped with something that limits the number of calls
749 (for example :func:`islice` or :func:`takewhile`). If not specified,
750 *fillvalue* defaults to ``None``.
749761e1
RH
751
752
116aa62b
GB
753.. _itertools-recipes:
754
1fa7682c
RH
755Itertools Recipes
756-----------------
116aa62b
GB
757
758This section shows recipes for creating an extended toolset using the existing
759itertools as building blocks.
760
f4ead487
RH
761The primary purpose of the itertools recipes is educational. The recipes show
762various ways of thinking about individual tools — for example, that
763``chain.from_iterable`` is related to the concept of flattening. The recipes
764also give ideas about ways that the tools can be combined — for example, how
c1c3be0f 765``compress()`` and ``range()`` can work together. The recipes also show patterns
f4ead487
RH
766for using itertools with the :mod:`operator` and :mod:`collections` modules as
767well as with the built-in itertools such as ``map()``, ``filter()``,
768``reversed()``, and ``enumerate()``.
769
770A secondary purpose of the recipes is to serve as an incubator. The
771``accumulate()``, ``compress()``, and ``pairwise()`` itertools started out as
f65fdbb8
RH
772recipes. Currently, the ``sliding_window()`` and ``iter_index()`` recipes
773are being tested to see whether they prove their worth.
f4ead487 774
adf02b36
RH
775Substantially all of these recipes and many, many others can be installed from
776the `more-itertools project <https://pypi.org/project/more-itertools/>`_ found
777on the Python Package Index::
778
dc14e33e 779 python -m pip install more-itertools
adf02b36 780
f4ead487
RH
781Many of the recipes offer the same high performance as the underlying toolset.
782Superior memory performance is kept by processing elements one at a time
116aa62b
GB
783rather than bringing the whole iterable into memory all at once. Code volume is
784kept small by linking the tools together in a functional style which helps
785eliminate temporary variables. High speed is retained by preferring
9afde1c0 786"vectorized" building blocks over the use of for-loops and :term:`generator`\s
fe337bfd
CH
787which incur interpreter overhead.
788
789.. testcode::
116aa62b 790
0769f957 791 import collections
c3453fbb 792 import functools
0769f957
RH
793 import math
794 import operator
795 import random
796
30c73624
RH
797 def take(n, iterable):
798 "Return first n items of the iterable as a list"
799 return list(islice(iterable, n))
800
9265dd72
RH
801 def prepend(value, iterator):
802 "Prepend a single value in front of an iterator"
8dc9b3fb 803 # prepend(1, [2, 3, 4]) --> 1 2 3 4
9265dd72
RH
804 return chain([value], iterator)
805
30c73624
RH
806 def tabulate(function, start=0):
807 "Return function(0), function(1), ..."
808 return map(function, count(start))
809
f65fdbb8
RH
810 def repeatfunc(func, times=None, *args):
811 """Repeat calls to func with specified arguments.
812
813 Example: repeatfunc(random.random)
814 """
815 if times is None:
816 return starmap(func, repeat(args))
817 return starmap(func, repeat(args, times))
818
819 def flatten(list_of_lists):
820 "Flatten one level of nesting"
821 return chain.from_iterable(list_of_lists)
822
823 def ncycles(iterable, n):
824 "Returns the sequence elements n times"
825 return chain.from_iterable(repeat(tuple(iterable), n))
826
dfe098d2
RH
827 def tail(n, iterable):
828 "Return an iterator over the last n items"
829 # tail(3, 'ABCDEFG') --> E F G
830 return iter(collections.deque(iterable, maxlen=n))
831
da1734c5
CS
832 def consume(iterator, n=None):
833 "Advance the iterator n-steps ahead. If n is None, consume entirely."
30c73624
RH
834 # Use functions that consume iterators at C speed.
835 if n is None:
836 # feed the entire iterator into a zero-length deque
837 collections.deque(iterator, maxlen=0)
838 else:
839 # advance to the empty slice starting at position n
840 next(islice(iterator, n, n), None)
841
842 def nth(iterable, n, default=None):
843 "Returns the nth item or a default value"
844 return next(islice(iterable, n, None), default)
845
846 def quantify(iterable, pred=bool):
f2636d2c 847 "Given a predicate that returns True or False, count the True results."
81bf10e4 848 "Count how many times the predicate is True"
30c73624
RH
849 return sum(map(pred, iterable))
850
f65fdbb8
RH
851 def all_equal(iterable):
852 "Returns True if all the elements are equal to each other"
853 g = groupby(iterable)
854 return next(g, True) and not next(g, False)
77fde8dc 855
f65fdbb8
RH
856 def first_true(iterable, default=False, pred=None):
857 """Returns the first true value in the iterable.
0214c7ad 858
f65fdbb8 859 If no true value is found, returns *default*
0214c7ad 860
f65fdbb8
RH
861 If *pred* is not None, returns the first item
862 for which pred(item) is true.
4075fe1d 863
4075fe1d 864 """
f65fdbb8
RH
865 # first_true([a,b,c], x) --> a or b or c or x
866 # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
867 return next(filter(pred, iterable), default)
0e28a3a5 868
f373c6b9 869 def iter_index(iterable, value, start=0, stop=None):
f4ead487 870 "Return indices where a value occurs in a sequence or iterable."
f4370318 871 # iter_index('AABCADEAF', 'A') --> 0 1 4 7
f373c6b9
RH
872 seq_index = getattr(iterable, 'index', None)
873 if seq_index is None:
f4ead487 874 # Slow path for general iterables
f373c6b9
RH
875 it = islice(iterable, start, stop)
876 for i, element in enumerate(it, start):
877 if element is value or element == value:
878 yield i
f4ead487
RH
879 else:
880 # Fast path for sequences
f2a55fec 881 stop = len(iterable) if stop is None else stop
f4ead487
RH
882 i = start - 1
883 try:
884 while True:
f373c6b9 885 yield (i := seq_index(value, i+1, stop))
f4ead487
RH
886 except ValueError:
887 pass
f4370318 888
f65fdbb8
RH
889 def iter_except(func, exception, first=None):
890 """ Call a function repeatedly until an exception is raised.
0769f957 891
f65fdbb8
RH
892 Converts a call-until-exception interface to an iterator interface.
893 Like builtins.iter(func, sentinel) but uses an exception instead
894 of a sentinel to end the loop.
30c73624 895
f65fdbb8
RH
896 Examples:
897 iter_except(functools.partial(heappop, h), IndexError) # priority queue iterator
898 iter_except(d.popitem, KeyError) # non-blocking dict iterator
899 iter_except(d.popleft, IndexError) # non-blocking deque iterator
900 iter_except(q.get_nowait, Queue.Empty) # loop over a producer Queue
901 iter_except(s.pop, KeyError) # non-blocking set iterator
30c73624 902
30c73624 903 """
f65fdbb8
RH
904 try:
905 if first is not None:
906 yield first() # For database APIs needing an initial cast to db.first()
907 while True:
908 yield func()
909 except exception:
910 pass
30c73624 911
270a0918 912 def grouper(iterable, n, *, incomplete='fill', fillvalue=None):
750368cb 913 "Collect data into non-overlapping fixed-length chunks or blocks"
270a0918
RH
914 # grouper('ABCDEFG', 3, fillvalue='x') --> ABC DEF Gxx
915 # grouper('ABCDEFG', 3, incomplete='strict') --> ABC DEF ValueError
916 # grouper('ABCDEFG', 3, incomplete='ignore') --> ABC DEF
30c73624 917 args = [iter(iterable)] * n
270a0918
RH
918 if incomplete == 'fill':
919 return zip_longest(*args, fillvalue=fillvalue)
920 if incomplete == 'strict':
921 return zip(*args, strict=True)
922 if incomplete == 'ignore':
923 return zip(*args)
924 else:
925 raise ValueError('Expected fill, strict, or ignore')
30c73624 926
750368cb 927 def sliding_window(iterable, n):
8dc9b3fb 928 # sliding_window('ABCDEFG', 4) --> ABCD BCDE CDEF DEFG
750368cb 929 it = iter(iterable)
423459be 930 window = collections.deque(islice(it, n-1), maxlen=n)
750368cb
RH
931 for x in it:
932 window.append(x)
933 yield tuple(window)
934
30c73624
RH
935 def roundrobin(*iterables):
936 "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
937 # Recipe credited to George Sakkis
337cbbac 938 num_active = len(iterables)
30c73624 939 nexts = cycle(iter(it).__next__ for it in iterables)
337cbbac 940 while num_active:
30c73624
RH
941 try:
942 for next in nexts:
943 yield next()
944 except StopIteration:
337cbbac
RH
945 # Remove the iterator we just exhausted from the cycle.
946 num_active -= 1
947 nexts = cycle(islice(nexts, num_active))
30c73624
RH
948
949 def partition(pred, iterable):
278030a1
RH
950 """Partition entries into false entries and true entries.
951
952 If *pred* is slow, consider wrapping it with functools.lru_cache().
953 """
30c73624
RH
954 # partition(is_odd, range(10)) --> 0 2 4 6 8 and 1 3 5 7 9
955 t1, t2 = tee(iterable)
956 return filterfalse(pred, t1), filter(pred, t2)
957
f65fdbb8
RH
958 def subslices(seq):
959 "Return all contiguous non-empty subslices of a sequence"
960 # subslices('ABCD') --> A AB ABC ABCD B BC BCD C CD D
961 slices = starmap(slice, combinations(range(len(seq) + 1), 2))
962 return map(operator.getitem, repeat(seq), slices)
963
91be41ad
RH
964 def before_and_after(predicate, it):
965 """ Variant of takewhile() that allows complete
966 access to the remainder of the iterator.
967
fcbf9b17
RH
968 >>> it = iter('ABCdEfGhI')
969 >>> all_upper, remainder = before_and_after(str.isupper, it)
970 >>> ''.join(all_upper)
91be41ad 971 'ABC'
fcbf9b17 972 >>> ''.join(remainder) # takewhile() would lose the 'd'
91be41ad
RH
973 'dEfGhI'
974
975 Note that the first iterator must be fully
976 consumed before the second iterator can
977 generate valid results.
978 """
979 it = iter(it)
980 transition = []
981 def true_iterator():
982 for elem in it:
983 if predicate(elem):
984 yield elem
985 else:
986 transition.append(elem)
987 return
988 def remainder_iterator():
989 yield from transition
990 yield from it
991 return true_iterator(), remainder_iterator()
992
30c73624
RH
993 def unique_everseen(iterable, key=None):
994 "List unique elements, preserving order. Remember all elements ever seen."
995 # unique_everseen('AAAABBBCCDAABBB') --> A B C D
4ebaae8a 996 # unique_everseen('ABBcCAD', str.lower) --> A B c D
30c73624 997 seen = set()
30c73624
RH
998 if key is None:
999 for element in filterfalse(seen.__contains__, iterable):
f4ead487 1000 seen.add(element)
30c73624 1001 yield element
4ebaae8a
RH
1002 # For order preserving deduplication,
1003 # a faster but non-lazy solution is:
f4ead487 1004 # yield from dict.fromkeys(iterable)
30c73624
RH
1005 else:
1006 for element in iterable:
1007 k = key(element)
1008 if k not in seen:
f4ead487 1009 seen.add(k)
30c73624 1010 yield element
4ebaae8a
RH
1011 # For use cases that allow the last matching element to be returned,
1012 # a faster but non-lazy solution is:
1013 # t1, t2 = tee(iterable)
1014 # yield from dict(zip(map(key, t1), t2)).values()
30c73624
RH
1015
1016 def unique_justseen(iterable, key=None):
1017 "List unique elements, preserving order. Remember only the element just seen."
1018 # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
4ebaae8a 1019 # unique_justseen('ABBcCAD', str.lower) --> A B c A D
b4c7f39b 1020 return map(next, map(operator.itemgetter(1), groupby(iterable, key)))
30c73624 1021
30c73624 1022
f65fdbb8 1023The following recipes have a more mathematical flavor:
30c73624 1024
f65fdbb8 1025.. testcode::
30c73624 1026
f65fdbb8
RH
1027 def powerset(iterable):
1028 "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1029 s = list(iterable)
1030 return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
1031
f65fdbb8
RH
1032 def sum_of_squares(it):
1033 "Add up the squares of the input values."
1034 # sum_of_squares([10, 20, 30]) -> 1400
1035 return math.sumprod(*tee(it))
31b26f63 1036
f65fdbb8
RH
1037 def transpose(it):
1038 "Swap the rows and columns of the input."
1039 # transpose([(1, 2, 3), (11, 22, 33)]) --> (1, 11) (2, 22) (3, 33)
1040 return zip(*it, strict=True)
31b26f63 1041
f65fdbb8
RH
1042 def matmul(m1, m2):
1043 "Multiply two matrices."
278030a1 1044 # matmul([(7, 5), (3, 5)], [(2, 5), (7, 9)]) --> (49, 80), (41, 60)
f65fdbb8
RH
1045 n = len(m2[0])
1046 return batched(starmap(math.sumprod, product(m1, transpose(m2))), n)
1047
1048 def convolve(signal, kernel):
f2636d2c
RH
1049 """Discrete linear convolution of two iterables.
1050
1051 The kernel is fully consumed before the calculations begin.
1052 The signal is consumed lazily and can be infinite.
1053
1054 Convolutions are mathematically commutative.
1055 If the signal and kernel are swapped,
1056 the output will be the same.
31b26f63 1057
f65fdbb8
RH
1058 Article: https://betterexplained.com/articles/intuitive-convolution/
1059 Video: https://www.youtube.com/watch?v=KuXjwB4LzSA
31b26f63 1060 """
f65fdbb8 1061 # convolve(data, [0.25, 0.25, 0.25, 0.25]) --> Moving average (blur)
f2636d2c
RH
1062 # convolve(data, [1/2, 0, -1/2]) --> 1st derivative estimate
1063 # convolve(data, [1, -2, 1]) --> 2nd derivative estimate
f65fdbb8
RH
1064 kernel = tuple(kernel)[::-1]
1065 n = len(kernel)
c3453fbb 1066 padded_signal = chain(repeat(0, n-1), signal, repeat(0, n-1))
77090370
RH
1067 windowed_signal = sliding_window(padded_signal, n)
1068 return map(math.sumprod, repeat(kernel), windowed_signal)
f65fdbb8
RH
1069
1070 def polynomial_from_roots(roots):
1071 """Compute a polynomial's coefficients from its roots.
1072
1073 (x - 5) (x + 4) (x - 3) expands to: x³ -4x² -17x + 60
1074 """
1075 # polynomial_from_roots([5, -4, 3]) --> [1, -4, -17, 60]
c3453fbb
RH
1076 factors = zip(repeat(1), map(operator.neg, roots))
1077 return list(functools.reduce(convolve, factors, [1]))
f65fdbb8
RH
1078
1079 def polynomial_eval(coefficients, x):
1080 """Evaluate a polynomial at a specific value.
1081
1082 Computes with better numeric stability than Horner's method.
1083 """
1084 # Evaluate x³ -4x² -17x + 60 at x = 2.5
1085 # polynomial_eval([1, -4, -17, 60], x=2.5) --> 8.125
1086 n = len(coefficients)
f2636d2c
RH
1087 if not n:
1088 return type(x)(0)
f65fdbb8
RH
1089 powers = map(pow, repeat(x), reversed(range(n)))
1090 return math.sumprod(coefficients, powers)
31b26f63 1091
278030a1
RH
1092 def polynomial_derivative(coefficients):
1093 """Compute the first derivative of a polynomial.
1094
1095 f(x) = x³ -4x² -17x + 60
1096 f'(x) = 3x² -8x -17
1097 """
1098 # polynomial_derivative([1, -4, -17, 60]) -> [3, -8, -17]
1099 n = len(coefficients)
1100 powers = reversed(range(1, n))
1101 return list(map(operator.mul, coefficients, powers))
1102
f2636d2c
RH
1103 def sieve(n):
1104 "Primes less than n."
1105 # sieve(30) --> 2 3 5 7 11 13 17 19 23 29
1106 if n > 2:
1107 yield 2
1108 start = 3
1109 data = bytearray((0, 1)) * (n // 2)
1110 limit = math.isqrt(n) + 1
1111 for p in iter_index(data, 1, start, limit):
1112 yield from iter_index(data, 1, start, p*p)
1113 data[p*p : n : p+p] = bytes(len(range(p*p, n, p+p)))
1114 start = p*p
1115 yield from iter_index(data, 1, start)
1116
1117 def factor(n):
1118 "Prime factors of n."
1119 # factor(99) --> 3 3 11
1120 # factor(1_000_000_000_000_007) --> 47 59 360620266859
1121 # factor(1_000_000_000_000_403) --> 1000000000000403
1122 for prime in sieve(math.isqrt(n) + 1):
1123 while not n % prime:
1124 yield prime
1125 n //= prime
1126 if n == 1:
1127 return
1128 if n > 1:
1129 yield n
1130
d37258dd 1131 def nth_combination(iterable, r, index):
fc40b302 1132 "Equivalent to list(combinations(iterable, r))[index]"
d37258dd
RH
1133 pool = tuple(iterable)
1134 n = len(pool)
59719a24 1135 c = math.comb(n, r)
d37258dd
RH
1136 if index < 0:
1137 index += c
1138 if index < 0 or index >= c:
1139 raise IndexError
1140 result = []
1141 while r:
1142 c, n, r = c*r//n, n-1, r-1
1143 while index >= c:
1144 index -= c
1145 c, n = c*(n-r)//n, n-1
1146 result.append(pool[-1-n])
1147 return tuple(result)
ee60550e 1148
f65fdbb8 1149
ee60550e
RH
1150.. doctest::
1151 :hide:
1152
1153 These examples no longer appear in the docs but are guaranteed
1154 to keep working.
1155
1156 >>> amounts = [120.15, 764.05, 823.14]
1157 >>> for checknum, amount in zip(count(1200), amounts):
1158 ... print('Check %d is for $%.2f' % (checknum, amount))
1159 ...
1160 Check 1200 is for $120.15
1161 Check 1201 is for $764.05
1162 Check 1202 is for $823.14
1163
1164 >>> import operator
1165 >>> for cube in map(operator.pow, range(1,4), repeat(3)):
1166 ... print(cube)
1167 ...
1168 1
1169 8
1170 27
1171
1172 >>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
1173 >>> for name in islice(reportlines, 3, None, 2):
1174 ... print(name.title())
1175 ...
1176 Alex
1177 Laura
1178 Martin
1179 Walter
1180 Samuele
1181
1182 >>> from operator import itemgetter
1183 >>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
1184 >>> di = sorted(sorted(d.items()), key=itemgetter(1))
1185 >>> for k, g in groupby(di, itemgetter(1)):
1186 ... print(k, list(map(itemgetter(0), g)))
1187 ...
1188 1 ['a', 'c', 'e']
1189 2 ['b', 'd', 'f']
1190 3 ['g']
1191
1192 # Find runs of consecutive numbers using groupby. The key to the solution
1193 # is differencing with a range so that consecutive numbers all appear in
1194 # same group.
1195 >>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
1196 >>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
1197 ... print(list(map(operator.itemgetter(1), g)))
1198 ...
1199 [1]
1200 [4, 5, 6]
1201 [10]
1202 [15, 16, 17, 18]
1203 [22]
1204 [25, 26, 27, 28]
1205
1206 Now, we test all of the itertool recipes
1207
ee60550e
RH
1208 >>> take(10, count())
1209 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1210
1211 >>> list(prepend(1, [2, 3, 4]))
1212 [1, 2, 3, 4]
1213
1214 >>> list(enumerate('abc'))
1215 [(0, 'a'), (1, 'b'), (2, 'c')]
1216
1217 >>> list(islice(tabulate(lambda x: 2*x), 4))
1218 [0, 2, 4, 6]
1219
1220 >>> list(tail(3, 'ABCDEFG'))
1221 ['E', 'F', 'G']
1222
1223 >>> it = iter(range(10))
1224 >>> consume(it, 3)
1225 >>> next(it)
1226 3
1227 >>> consume(it)
1228 >>> next(it, 'Done')
1229 'Done'
1230
1231 >>> nth('abcde', 3)
1232 'd'
1233
1234 >>> nth('abcde', 9) is None
1235 True
1236
1237 >>> [all_equal(s) for s in ('', 'A', 'AAAA', 'AAAB', 'AAABA')]
1238 [True, True, True, False, False]
1239
1240 >>> quantify(range(99), lambda x: x%2==0)
1241 50
1242
1243 >>> quantify([True, False, False, True, True])
1244 3
1245
1246 >>> quantify(range(12), pred=lambda x: x%2==1)
1247 6
1248
1249 >>> a = [[1, 2, 3], [4, 5, 6]]
1250 >>> list(flatten(a))
1251 [1, 2, 3, 4, 5, 6]
1252
1253 >>> list(repeatfunc(pow, 5, 2, 3))
1254 [8, 8, 8, 8, 8]
1255
ee60550e
RH
1256 >>> take(5, map(int, repeatfunc(random.random)))
1257 [0, 0, 0, 0, 0]
1258
ee60550e
RH
1259 >>> list(ncycles('abc', 3))
1260 ['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1261
6bde3d2f
RH
1262 >>> sum_of_squares([10, 20, 30])
1263 1400
1264
1265 >>> list(transpose([(1, 2, 3), (11, 22, 33)]))
1266 [(1, 11), (2, 22), (3, 33)]
1267
1268 >>> list(matmul([(7, 5), (3, 5)], [[2, 5], [7, 9]]))
1269 [(49, 80), (41, 60)]
1270 >>> list(matmul([[2, 5], [7, 9], [3, 4]], [[7, 11, 5, 4, 9], [3, 5, 2, 6, 3]]))
1271 [(29, 47, 20, 38, 33), (76, 122, 53, 82, 90), (33, 53, 23, 36, 39)]
1272
ee60550e
RH
1273 >>> data = [20, 40, 24, 32, 20, 28, 16]
1274 >>> list(convolve(data, [0.25, 0.25, 0.25, 0.25]))
1275 [5.0, 15.0, 21.0, 29.0, 29.0, 26.0, 24.0, 16.0, 11.0, 4.0]
1276 >>> list(convolve(data, [1, -1]))
1277 [20, 20, -16, 8, -12, 8, -12, -16]
1278 >>> list(convolve(data, [1, -2, 1]))
1279 [20, 0, -36, 24, -20, 20, -20, -4, 16]
1280
094cf392
RH
1281 >>> from fractions import Fraction
1282 >>> from decimal import Decimal
1283 >>> polynomial_eval([1, -4, -17, 60], x=2)
1284 18
1285 >>> x = 2; x**3 - 4*x**2 -17*x + 60
1286 18
1287 >>> polynomial_eval([1, -4, -17, 60], x=2.5)
1288 8.125
1289 >>> x = 2.5; x**3 - 4*x**2 -17*x + 60
1290 8.125
1291 >>> polynomial_eval([1, -4, -17, 60], x=Fraction(2, 3))
1292 Fraction(1274, 27)
1293 >>> x = Fraction(2, 3); x**3 - 4*x**2 -17*x + 60
1294 Fraction(1274, 27)
1295 >>> polynomial_eval([1, -4, -17, 60], x=Decimal('1.75'))
1296 Decimal('23.359375')
1297 >>> x = Decimal('1.75'); x**3 - 4*x**2 -17*x + 60
1298 Decimal('23.359375')
1299 >>> polynomial_eval([], 2)
1300 0
1301 >>> polynomial_eval([], 2.5)
1302 0.0
1303 >>> polynomial_eval([], Fraction(2, 3))
1304 Fraction(0, 1)
1305 >>> polynomial_eval([], Decimal('1.75'))
f2636d2c 1306 Decimal('0')
094cf392
RH
1307 >>> polynomial_eval([11], 7) == 11
1308 True
1309 >>> polynomial_eval([11, 2], 7) == 11 * 7 + 2
1310 True
1311
0e28a3a5
RH
1312 >>> polynomial_from_roots([5, -4, 3])
1313 [1, -4, -17, 60]
1314 >>> factored = lambda x: (x - 5) * (x + 4) * (x - 3)
1315 >>> expanded = lambda x: x**3 -4*x**2 -17*x + 60
1316 >>> all(factored(x) == expanded(x) for x in range(-10, 11))
1317 True
1318
278030a1
RH
1319 >>> polynomial_derivative([1, -4, -17, 60])
1320 [3, -8, -17]
1321
f4370318
RH
1322 >>> list(iter_index('AABCADEAF', 'A'))
1323 [0, 1, 4, 7]
1324 >>> list(iter_index('AABCADEAF', 'B'))
1325 [2]
1326 >>> list(iter_index('AABCADEAF', 'X'))
1327 []
1328 >>> list(iter_index('', 'X'))
1329 []
f4ead487
RH
1330 >>> list(iter_index('AABCADEAF', 'A', 1))
1331 [1, 4, 7]
1332 >>> list(iter_index(iter('AABCADEAF'), 'A', 1))
1333 [1, 4, 7]
1334 >>> list(iter_index('AABCADEAF', 'A', 2))
1335 [4, 7]
1336 >>> list(iter_index(iter('AABCADEAF'), 'A', 2))
1337 [4, 7]
1338 >>> list(iter_index('AABCADEAF', 'A', 10))
1339 []
1340 >>> list(iter_index(iter('AABCADEAF'), 'A', 10))
1341 []
f373c6b9
RH
1342 >>> list(iter_index('AABCADEAF', 'A', 1, 7))
1343 [1, 4]
1344 >>> list(iter_index(iter('AABCADEAF'), 'A', 1, 7))
1345 [1, 4]
1346 >>> # Verify that ValueErrors not swallowed (gh-107208)
1347 >>> def assert_no_value(iterable, forbidden_value):
1348 ... for item in iterable:
1349 ... if item == forbidden_value:
1350 ... raise ValueError
1351 ... yield item
1352 ...
1353 >>> list(iter_index(assert_no_value('AABCADEAF', 'B'), 'A'))
1354 Traceback (most recent call last):
1355 ...
1356 ValueError
f2a55fec
RH
1357 >>> # Verify that both paths can find identical NaN values
1358 >>> x = float('NaN')
1359 >>> y = float('NaN')
1360 >>> list(iter_index([0, x, x, y, 0], x))
1361 [1, 2]
1362 >>> list(iter_index(iter([0, x, x, y, 0]), x))
1363 [1, 2]
1364 >>> # Test list input. Lists do not support None for the stop argument
1365 >>> list(iter_index(list('AABCADEAF'), 'A'))
1366 [0, 1, 4, 7]
f4370318 1367
8dc9b3fb
RH
1368 >>> list(sieve(30))
1369 [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]
e500cc04
RH
1370 >>> small_primes = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97]
1371 >>> all(list(sieve(n)) == [p for p in small_primes if p < n] for n in range(101))
78359b1d 1372 True
8dc9b3fb
RH
1373 >>> len(list(sieve(100)))
1374 25
1375 >>> len(list(sieve(1_000)))
1376 168
1377 >>> len(list(sieve(10_000)))
1378 1229
1379 >>> len(list(sieve(100_000)))
1380 9592
1381 >>> len(list(sieve(1_000_000)))
1382 78498
78359b1d
RH
1383 >>> carmichael = {561, 1105, 1729, 2465, 2821, 6601, 8911} # https://oeis.org/A002997
1384 >>> set(sieve(10_000)).isdisjoint(carmichael)
1385 True
8dc9b3fb 1386
babb22da
RH
1387 >>> list(factor(99)) # Code example 1
1388 [3, 3, 11]
1389 >>> list(factor(1_000_000_000_000_007)) # Code example 2
1390 [47, 59, 360620266859]
1391 >>> list(factor(1_000_000_000_000_403)) # Code example 3
1392 [1000000000000403]
2d524068 1393 >>> list(factor(0))
0769f957 1394 []
2d524068 1395 >>> list(factor(1))
0769f957 1396 []
2d524068 1397 >>> list(factor(2))
0769f957 1398 [2]
2d524068 1399 >>> list(factor(3))
0769f957 1400 [3]
2d524068 1401 >>> list(factor(4))
0769f957 1402 [2, 2]
2d524068 1403 >>> list(factor(5))
0769f957 1404 [5]
2d524068 1405 >>> list(factor(6))
0769f957 1406 [2, 3]
2d524068 1407 >>> list(factor(7))
0769f957 1408 [7]
2d524068 1409 >>> list(factor(8))
0769f957 1410 [2, 2, 2]
2d524068 1411 >>> list(factor(9))
0769f957 1412 [3, 3]
2d524068 1413 >>> list(factor(10))
0769f957 1414 [2, 5]
c4c57901
RH
1415 >>> list(factor(128_884_753_939)) # large prime
1416 [128884753939]
1417 >>> list(factor(999953 * 999983)) # large semiprime
2d524068 1418 [999953, 999983]
c4c57901
RH
1419 >>> list(factor(6 ** 20)) == [2] * 20 + [3] * 20 # large power
1420 True
1421 >>> list(factor(909_909_090_909)) # large multiterm composite
1422 [3, 3, 7, 13, 13, 751, 113797]
1423 >>> math.prod([3, 3, 7, 13, 13, 751, 113797])
1424 909909090909
1425 >>> all(math.prod(factor(n)) == n for n in range(1, 2_000))
0769f957 1426 True
c4c57901 1427 >>> all(set(factor(n)) <= set(sieve(n+1)) for n in range(2_000))
0769f957 1428 True
c4c57901 1429 >>> all(list(factor(n)) == sorted(factor(n)) for n in range(2_000))
0769f957
RH
1430 True
1431
ee60550e
RH
1432 >>> list(flatten([('a', 'b'), (), ('c', 'd', 'e'), ('f',), ('g', 'h', 'i')]))
1433 ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i']
1434
ee60550e
RH
1435 >>> random.seed(85753098575309)
1436 >>> list(repeatfunc(random.random, 3))
1437 [0.16370491282496968, 0.45889608687313455, 0.3747076837820118]
1438 >>> list(repeatfunc(chr, 3, 65))
1439 ['A', 'A', 'A']
1440 >>> list(repeatfunc(pow, 3, 2, 5))
1441 [32, 32, 32]
1442
1443 >>> list(grouper('abcdefg', 3, fillvalue='x'))
1444 [('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1445
1446 >>> it = grouper('abcdefg', 3, incomplete='strict')
1447 >>> next(it)
1448 ('a', 'b', 'c')
1449 >>> next(it)
1450 ('d', 'e', 'f')
1451 >>> next(it)
1452 Traceback (most recent call last):
1453 ...
1454 ValueError: zip() argument 2 is shorter than argument 1
1455
1456 >>> list(grouper('abcdefg', n=3, incomplete='ignore'))
1457 [('a', 'b', 'c'), ('d', 'e', 'f')]
1458
423459be
RH
1459 >>> list(sliding_window('ABCDEFG', 1))
1460 [('A',), ('B',), ('C',), ('D',), ('E',), ('F',), ('G',)]
1461 >>> list(sliding_window('ABCDEFG', 2))
1462 [('A', 'B'), ('B', 'C'), ('C', 'D'), ('D', 'E'), ('E', 'F'), ('F', 'G')]
1463 >>> list(sliding_window('ABCDEFG', 3))
1464 [('A', 'B', 'C'), ('B', 'C', 'D'), ('C', 'D', 'E'), ('D', 'E', 'F'), ('E', 'F', 'G')]
ee60550e
RH
1465 >>> list(sliding_window('ABCDEFG', 4))
1466 [('A', 'B', 'C', 'D'), ('B', 'C', 'D', 'E'), ('C', 'D', 'E', 'F'), ('D', 'E', 'F', 'G')]
423459be
RH
1467 >>> list(sliding_window('ABCDEFG', 5))
1468 [('A', 'B', 'C', 'D', 'E'), ('B', 'C', 'D', 'E', 'F'), ('C', 'D', 'E', 'F', 'G')]
1469 >>> list(sliding_window('ABCDEFG', 6))
1470 [('A', 'B', 'C', 'D', 'E', 'F'), ('B', 'C', 'D', 'E', 'F', 'G')]
1471 >>> list(sliding_window('ABCDEFG', 7))
1472 [('A', 'B', 'C', 'D', 'E', 'F', 'G')]
1473 >>> list(sliding_window('ABCDEFG', 8))
1474 []
1475 >>> try:
1476 ... list(sliding_window('ABCDEFG', -1))
1477 ... except ValueError:
1478 ... 'zero or negative n not supported'
1479 ...
1480 'zero or negative n not supported'
1481 >>> try:
1482 ... list(sliding_window('ABCDEFG', 0))
1483 ... except ValueError:
1484 ... 'zero or negative n not supported'
1485 ...
1486 'zero or negative n not supported'
ee60550e
RH
1487
1488 >>> list(roundrobin('abc', 'd', 'ef'))
1489 ['a', 'd', 'e', 'b', 'f', 'c']
1490
1491 >>> def is_odd(x):
1492 ... return x % 2 == 1
1493
1494 >>> evens, odds = partition(is_odd, range(10))
1495 >>> list(evens)
1496 [0, 2, 4, 6, 8]
1497 >>> list(odds)
1498 [1, 3, 5, 7, 9]
1499
1500 >>> it = iter('ABCdEfGhI')
1501 >>> all_upper, remainder = before_and_after(str.isupper, it)
1502 >>> ''.join(all_upper)
1503 'ABC'
1504 >>> ''.join(remainder)
1505 'dEfGhI'
1506
06a49117
RH
1507 >>> list(subslices('ABCD'))
1508 ['A', 'AB', 'ABC', 'ABCD', 'B', 'BC', 'BCD', 'C', 'CD', 'D']
1509
ee60550e
RH
1510 >>> list(powerset([1,2,3]))
1511 [(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
1512
1513 >>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1514 True
1515
1516 >>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1517 True
1518
1519 >>> list(unique_everseen('AAAABBBCCDAABBB'))
1520 ['A', 'B', 'C', 'D']
ee60550e
RH
1521 >>> list(unique_everseen('ABBCcAD', str.lower))
1522 ['A', 'B', 'C', 'D']
4ebaae8a
RH
1523 >>> list(unique_everseen('ABBcCAD', str.lower))
1524 ['A', 'B', 'c', 'D']
ee60550e
RH
1525
1526 >>> list(unique_justseen('AAAABBBCCDAABBB'))
1527 ['A', 'B', 'C', 'D', 'A', 'B']
ee60550e
RH
1528 >>> list(unique_justseen('ABBCcAD', str.lower))
1529 ['A', 'B', 'C', 'A', 'D']
4ebaae8a
RH
1530 >>> list(unique_justseen('ABBcCAD', str.lower))
1531 ['A', 'B', 'c', 'A', 'D']
ee60550e
RH
1532
1533 >>> d = dict(a=1, b=2, c=3)
1534 >>> it = iter_except(d.popitem, KeyError)
1535 >>> d['d'] = 4
1536 >>> next(it)
1537 ('d', 4)
1538 >>> next(it)
1539 ('c', 3)
1540 >>> next(it)
1541 ('b', 2)
1542 >>> d['e'] = 5
1543 >>> next(it)
1544 ('e', 5)
1545 >>> next(it)
1546 ('a', 1)
1547 >>> next(it, 'empty')
1548 'empty'
1549
1550 >>> first_true('ABC0DEF1', '9', str.isdigit)
1551 '0'
1552
1553 >>> population = 'ABCDEFGH'
1554 >>> for r in range(len(population) + 1):
1555 ... seq = list(combinations(population, r))
1556 ... for i in range(len(seq)):
1557 ... assert nth_combination(population, r, i) == seq[i]
1558 ... for i in range(-len(seq), 0):
1559 ... assert nth_combination(population, r, i) == seq[i]
1560
1561 >>> iterable = 'abcde'
1562 >>> r = 3
1563 >>> combos = list(combinations(iterable, r))
1564 >>> all(nth_combination(iterable, r, i) == comb for i, comb in enumerate(combos))
1565 True
f65fdbb8
RH
1566
1567
1568.. testcode::
1569 :hide:
1570
1571 # Old recipes and their tests which are guaranteed to continue to work.
1572
1573 def sumprod(vec1, vec2):
1574 "Compute a sum of products."
1575 return sum(starmap(operator.mul, zip(vec1, vec2, strict=True)))
1576
1577 def dotproduct(vec1, vec2):
1578 return sum(map(operator.mul, vec1, vec2))
1579
1580 def pad_none(iterable):
1581 """Returns the sequence elements and then returns None indefinitely.
1582
1583 Useful for emulating the behavior of the built-in map() function.
1584 """
1585 return chain(iterable, repeat(None))
1586
1587 def triplewise(iterable):
1588 "Return overlapping triplets from an iterable"
1589 # triplewise('ABCDEFG') --> ABC BCD CDE DEF EFG
1590 for (a, _), (b, c) in pairwise(pairwise(iterable)):
1591 yield a, b, c
1592
1593
1594.. doctest::
1595 :hide:
1596
1597 >>> dotproduct([1,2,3], [4,5,6])
1598 32
1599
1600 >>> sumprod([1,2,3], [4,5,6])
1601 32
1602
1603 >>> list(islice(pad_none('abc'), 0, 6))
1604 ['a', 'b', 'c', None, None, None]
1605
1606 >>> list(triplewise('ABCDEFG'))
1607 [('A', 'B', 'C'), ('B', 'C', 'D'), ('C', 'D', 'E'), ('D', 'E', 'F'), ('E', 'F', 'G')]