]>
Commit | Line | Data |
---|---|---|
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 |
20 | This module implements a number of :term:`iterator` building blocks inspired |
21 | by constructs from APL, Haskell, and SML. Each has been recast in a form | |
22 | suitable for Python. | |
116aa62b GB |
23 | |
24 | The module standardizes a core set of fast, memory efficient tools that are | |
f76b9209 RH |
25 | useful by themselves or in combination. Together, they form an "iterator |
26 | algebra" making it possible to construct specialized tools succinctly and | |
27 | efficiently in pure Python. | |
116aa62b GB |
28 | |
29 | For instance, SML provides a tabulation tool: ``tabulate(f)`` which produces a | |
b660599f | 30 | sequence ``f(0), f(1), ...``. The same effect can be achieved in Python |
a6c6037f | 31 | by combining :func:`map` and :func:`count` to form ``map(f, count())``. |
116aa62b | 32 | |
2c109ab5 RH |
33 | These tools and their built-in counterparts also work well with the high-speed |
34 | functions in the :mod:`operator` module. For example, the multiplication | |
35 | operator 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 | ================== ================= ================================================= ========================================= |
42 | Iterator 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 | ============================ ============================ ================================================= ============================================================= |
52 | Iterator 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 | ============================================== ==================== ============================================================= |
73 | Iterator 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 | ============================================== ============================================================= |
82 | Examples 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 | ||
93 | Itertool functions | |
94 | ------------------ | |
95 | ||
96 | The following module functions all construct and return iterators. Some provide | |
97 | streams of infinite length, so they should only be accessed by functions or | |
98 | loops 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 |
755 | Itertools Recipes |
756 | ----------------- | |
116aa62b GB |
757 | |
758 | This section shows recipes for creating an extended toolset using the existing | |
759 | itertools as building blocks. | |
760 | ||
f4ead487 RH |
761 | The primary purpose of the itertools recipes is educational. The recipes show |
762 | various ways of thinking about individual tools — for example, that | |
763 | ``chain.from_iterable`` is related to the concept of flattening. The recipes | |
764 | also 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 |
766 | for using itertools with the :mod:`operator` and :mod:`collections` modules as |
767 | well as with the built-in itertools such as ``map()``, ``filter()``, | |
768 | ``reversed()``, and ``enumerate()``. | |
769 | ||
770 | A secondary purpose of the recipes is to serve as an incubator. The | |
771 | ``accumulate()``, ``compress()``, and ``pairwise()`` itertools started out as | |
f65fdbb8 RH |
772 | recipes. Currently, the ``sliding_window()`` and ``iter_index()`` recipes |
773 | are being tested to see whether they prove their worth. | |
f4ead487 | 774 | |
adf02b36 RH |
775 | Substantially all of these recipes and many, many others can be installed from |
776 | the `more-itertools project <https://pypi.org/project/more-itertools/>`_ found | |
777 | on the Python Package Index:: | |
778 | ||
dc14e33e | 779 | python -m pip install more-itertools |
adf02b36 | 780 | |
f4ead487 RH |
781 | Many of the recipes offer the same high performance as the underlying toolset. |
782 | Superior memory performance is kept by processing elements one at a time | |
116aa62b GB |
783 | rather than bringing the whole iterable into memory all at once. Code volume is |
784 | kept small by linking the tools together in a functional style which helps | |
785 | eliminate 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 |
787 | which 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 | 1023 | The 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')] |