Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 1 | :mod:`itertools` --- Functions creating iterators for efficient looping |
| 2 | ======================================================================= |
| 3 | |
| 4 | .. module:: itertools |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 5 | :synopsis: Functions creating iterators for efficient looping. |
Terry Jan Reedy | fa089b9 | 2016-06-11 15:02:54 -0400 | [diff] [blame] | 6 | |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 7 | .. moduleauthor:: Raymond Hettinger <python@rcn.com> |
| 8 | .. sectionauthor:: Raymond Hettinger <python@rcn.com> |
| 9 | |
Christian Heimes | fe337bf | 2008-03-23 21:54:12 +0000 | [diff] [blame] | 10 | .. testsetup:: |
| 11 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 12 | from itertools import * |
Christian Heimes | fe337bf | 2008-03-23 21:54:12 +0000 | [diff] [blame] | 13 | |
Terry Jan Reedy | fa089b9 | 2016-06-11 15:02:54 -0400 | [diff] [blame] | 14 | -------------- |
Christian Heimes | fe337bf | 2008-03-23 21:54:12 +0000 | [diff] [blame] | 15 | |
Raymond Hettinger | f76b920 | 2009-02-17 20:00:59 +0000 | [diff] [blame] | 16 | This module implements a number of :term:`iterator` building blocks inspired |
| 17 | by constructs from APL, Haskell, and SML. Each has been recast in a form |
| 18 | suitable for Python. |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 19 | |
| 20 | The module standardizes a core set of fast, memory efficient tools that are |
Raymond Hettinger | f76b920 | 2009-02-17 20:00:59 +0000 | [diff] [blame] | 21 | useful by themselves or in combination. Together, they form an "iterator |
| 22 | algebra" making it possible to construct specialized tools succinctly and |
| 23 | efficiently in pure Python. |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 24 | |
| 25 | For instance, SML provides a tabulation tool: ``tabulate(f)`` which produces a |
Ezio Melotti | b660599 | 2010-01-21 20:57:24 +0000 | [diff] [blame] | 26 | sequence ``f(0), f(1), ...``. The same effect can be achieved in Python |
Raymond Hettinger | a6c6037 | 2008-03-13 01:26:19 +0000 | [diff] [blame] | 27 | by combining :func:`map` and :func:`count` to form ``map(f, count())``. |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 28 | |
Raymond Hettinger | 2c109ab | 2009-03-12 00:29:44 +0000 | [diff] [blame] | 29 | These tools and their built-in counterparts also work well with the high-speed |
| 30 | functions in the :mod:`operator` module. For example, the multiplication |
| 31 | operator can be mapped across two vectors to form an efficient dot-product: |
| 32 | ``sum(map(operator.mul, vector1, vector2))``. |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 33 | |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 34 | |
Raymond Hettinger | 6693d7a | 2017-12-15 13:17:55 -0800 | [diff] [blame] | 35 | **Infinite iterators:** |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 36 | |
Raymond Hettinger | 5bfd8ce | 2009-04-10 19:02:36 +0000 | [diff] [blame] | 37 | ================== ================= ================================================= ========================================= |
| 38 | Iterator Arguments Results Example |
| 39 | ================== ================= ================================================= ========================================= |
| 40 | :func:`count` start, [step] start, start+step, start+2*step, ... ``count(10) --> 10 11 12 13 14 ...`` |
| 41 | :func:`cycle` p p0, p1, ... plast, p0, p1, ... ``cycle('ABCD') --> A B C D A B C D ...`` |
| 42 | :func:`repeat` elem [,n] elem, elem, elem, ... endlessly or up to n times ``repeat(10, 3) --> 10 10 10`` |
| 43 | ================== ================= ================================================= ========================================= |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 44 | |
Raymond Hettinger | f76b920 | 2009-02-17 20:00:59 +0000 | [diff] [blame] | 45 | **Iterators terminating on the shortest input sequence:** |
| 46 | |
Benjamin Peterson | 2989f58 | 2014-01-16 10:10:13 -0500 | [diff] [blame] | 47 | ============================ ============================ ================================================= ============================================================= |
| 48 | Iterator Arguments Results Example |
| 49 | ============================ ============================ ================================================= ============================================================= |
| 50 | :func:`accumulate` p [,func] p0, p0+p1, p0+p1+p2, ... ``accumulate([1,2,3,4,5]) --> 1 3 6 10 15`` |
| 51 | :func:`chain` p, q, ... p0, p1, ... plast, q0, q1, ... ``chain('ABC', 'DEF') --> A B C D E F`` |
| 52 | :func:`chain.from_iterable` iterable p0, p1, ... plast, q0, q1, ... ``chain.from_iterable(['ABC', 'DEF']) --> A B C D E F`` |
| 53 | :func:`compress` data, selectors (d[0] if s[0]), (d[1] if s[1]), ... ``compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F`` |
| 54 | :func:`dropwhile` pred, seq seq[n], seq[n+1], starting when pred fails ``dropwhile(lambda x: x<5, [1,4,6,4,1]) --> 6 4 1`` |
| 55 | :func:`filterfalse` pred, seq elements of seq where pred(elem) is false ``filterfalse(lambda x: x%2, range(10)) --> 0 2 4 6 8`` |
Raymond Hettinger | 49392c6 | 2017-09-25 01:21:06 -0700 | [diff] [blame] | 56 | :func:`groupby` iterable[, key] sub-iterators grouped by value of key(v) |
Benjamin Peterson | 2989f58 | 2014-01-16 10:10:13 -0500 | [diff] [blame] | 57 | :func:`islice` seq, [start,] stop [, step] elements from seq[start:stop:step] ``islice('ABCDEFG', 2, None) --> C D E F G`` |
Raymond Hettinger | cc061d0 | 2020-11-30 20:42:54 -0800 | [diff] [blame] | 58 | :func:`pairwise` iterable (p[0], p[1]), (p[1], p[2]) ``pairwise('ABCDEFG') --> AB BC CD DE EF FG`` |
Benjamin Peterson | 2989f58 | 2014-01-16 10:10:13 -0500 | [diff] [blame] | 59 | :func:`starmap` func, seq func(\*seq[0]), func(\*seq[1]), ... ``starmap(pow, [(2,5), (3,2), (10,3)]) --> 32 9 1000`` |
| 60 | :func:`takewhile` pred, seq seq[0], seq[1], until pred fails ``takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4`` |
| 61 | :func:`tee` it, n it1, it2, ... itn splits one iterator into n |
| 62 | :func:`zip_longest` p, q, ... (p[0], q[0]), (p[1], q[1]), ... ``zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-`` |
| 63 | ============================ ============================ ================================================= ============================================================= |
Raymond Hettinger | f76b920 | 2009-02-17 20:00:59 +0000 | [diff] [blame] | 64 | |
Raymond Hettinger | 6693d7a | 2017-12-15 13:17:55 -0800 | [diff] [blame] | 65 | **Combinatoric iterators:** |
Raymond Hettinger | f76b920 | 2009-02-17 20:00:59 +0000 | [diff] [blame] | 66 | |
Raymond Hettinger | 7f587cd | 2009-04-10 19:43:50 +0000 | [diff] [blame] | 67 | ============================================== ==================== ============================================================= |
| 68 | Iterator Arguments Results |
| 69 | ============================================== ==================== ============================================================= |
| 70 | :func:`product` p, q, ... [repeat=1] cartesian product, equivalent to a nested for-loop |
| 71 | :func:`permutations` p[, r] r-length tuples, all possible orderings, no repeated elements |
Raymond Hettinger | 36c3c02 | 2009-11-19 01:20:07 +0000 | [diff] [blame] | 72 | :func:`combinations` p, r r-length tuples, in sorted order, no repeated elements |
| 73 | :func:`combinations_with_replacement` p, r r-length tuples, in sorted order, with repeated elements |
Raymond Hettinger | 7f587cd | 2009-04-10 19:43:50 +0000 | [diff] [blame] | 74 | ============================================== ==================== ============================================================= |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 75 | |
Benedikt Werner | 14e3c44 | 2019-03-21 16:28:49 +0100 | [diff] [blame] | 76 | ============================================== ============================================================= |
| 77 | Examples Results |
| 78 | ============================================== ============================================================= |
| 79 | ``product('ABCD', repeat=2)`` ``AA AB AC AD BA BB BC BD CA CB CC CD DA DB DC DD`` |
| 80 | ``permutations('ABCD', 2)`` ``AB AC AD BA BC BD CA CB CD DA DB DC`` |
| 81 | ``combinations('ABCD', 2)`` ``AB AC AD BC BD CD`` |
| 82 | ``combinations_with_replacement('ABCD', 2)`` ``AA AB AC AD BB BC BD CC CD DD`` |
| 83 | ============================================== ============================================================= |
| 84 | |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 85 | |
| 86 | .. _itertools-functions: |
| 87 | |
| 88 | Itertool functions |
| 89 | ------------------ |
| 90 | |
| 91 | The following module functions all construct and return iterators. Some provide |
| 92 | streams of infinite length, so they should only be accessed by functions or |
| 93 | loops that truncate the stream. |
| 94 | |
Lisa Roach | 9718b59 | 2018-09-23 17:34:59 -0700 | [diff] [blame] | 95 | .. function:: accumulate(iterable[, func, *, initial=None]) |
Raymond Hettinger | adb8146 | 2010-12-01 22:50:36 +0000 | [diff] [blame] | 96 | |
Andrew Kuchling | 15b04eb | 2014-04-15 22:28:40 -0400 | [diff] [blame] | 97 | Make an iterator that returns accumulated sums, or accumulated |
| 98 | results of other binary functions (specified via the optional |
Lisa Roach | 9718b59 | 2018-09-23 17:34:59 -0700 | [diff] [blame] | 99 | *func* argument). |
| 100 | |
| 101 | If *func* is supplied, it should be a function |
Andrew Kuchling | 15b04eb | 2014-04-15 22:28:40 -0400 | [diff] [blame] | 102 | of two arguments. Elements of the input *iterable* may be any type |
| 103 | that can be accepted as arguments to *func*. (For example, with |
| 104 | the default operation of addition, elements may be any addable |
| 105 | type including :class:`~decimal.Decimal` or |
Lisa Roach | 9718b59 | 2018-09-23 17:34:59 -0700 | [diff] [blame] | 106 | :class:`~fractions.Fraction`.) |
| 107 | |
| 108 | Usually, the number of elements output matches the input iterable. |
| 109 | However, if the keyword argument *initial* is provided, the |
| 110 | accumulation leads off with the *initial* value so that the output |
| 111 | has one more element than the input iterable. |
Raymond Hettinger | adb8146 | 2010-12-01 22:50:36 +0000 | [diff] [blame] | 112 | |
Raymond Hettinger | 672866d | 2016-05-28 00:17:54 -0700 | [diff] [blame] | 113 | Roughly equivalent to:: |
Raymond Hettinger | 5d44613 | 2011-03-27 18:52:10 -0700 | [diff] [blame] | 114 | |
Lisa Roach | 9718b59 | 2018-09-23 17:34:59 -0700 | [diff] [blame] | 115 | def accumulate(iterable, func=operator.add, *, initial=None): |
Raymond Hettinger | adb8146 | 2010-12-01 22:50:36 +0000 | [diff] [blame] | 116 | 'Return running totals' |
Raymond Hettinger | d8ff465 | 2010-12-03 02:09:34 +0000 | [diff] [blame] | 117 | # accumulate([1,2,3,4,5]) --> 1 3 6 10 15 |
Lisa Roach | 9718b59 | 2018-09-23 17:34:59 -0700 | [diff] [blame] | 118 | # accumulate([1,2,3,4,5], initial=100) --> 100 101 103 106 110 115 |
Raymond Hettinger | 5d44613 | 2011-03-27 18:52:10 -0700 | [diff] [blame] | 119 | # accumulate([1,2,3,4,5], operator.mul) --> 1 2 6 24 120 |
Raymond Hettinger | d8ff465 | 2010-12-03 02:09:34 +0000 | [diff] [blame] | 120 | it = iter(iterable) |
Lisa Roach | 9718b59 | 2018-09-23 17:34:59 -0700 | [diff] [blame] | 121 | total = initial |
| 122 | if initial is None: |
| 123 | try: |
| 124 | total = next(it) |
| 125 | except StopIteration: |
| 126 | return |
Raymond Hettinger | d8ff465 | 2010-12-03 02:09:34 +0000 | [diff] [blame] | 127 | yield total |
| 128 | for element in it: |
Raymond Hettinger | 5d44613 | 2011-03-27 18:52:10 -0700 | [diff] [blame] | 129 | total = func(total, element) |
Raymond Hettinger | d8ff465 | 2010-12-03 02:09:34 +0000 | [diff] [blame] | 130 | yield total |
Raymond Hettinger | adb8146 | 2010-12-01 22:50:36 +0000 | [diff] [blame] | 131 | |
Raymond Hettinger | 295c1d4 | 2011-04-21 11:09:28 -0700 | [diff] [blame] | 132 | There are a number of uses for the *func* argument. It can be set to |
| 133 | :func:`min` for a running minimum, :func:`max` for a running maximum, or |
| 134 | :func:`operator.mul` for a running product. Amortization tables can be |
| 135 | built by accumulating interest and applying payments. First-order |
Georg Brandl | 5d94134 | 2016-02-26 19:37:12 +0100 | [diff] [blame] | 136 | `recurrence relations <https://en.wikipedia.org/wiki/Recurrence_relation>`_ |
Raymond Hettinger | 295c1d4 | 2011-04-21 11:09:28 -0700 | [diff] [blame] | 137 | can be modeled by supplying the initial value in the iterable and using only |
| 138 | the accumulated total in *func* argument:: |
Raymond Hettinger | 5d44613 | 2011-03-27 18:52:10 -0700 | [diff] [blame] | 139 | |
| 140 | >>> data = [3, 4, 6, 2, 1, 9, 0, 7, 5, 8] |
| 141 | >>> list(accumulate(data, operator.mul)) # running product |
| 142 | [3, 12, 72, 144, 144, 1296, 0, 0, 0, 0] |
| 143 | >>> list(accumulate(data, max)) # running maximum |
| 144 | [3, 4, 6, 6, 6, 9, 9, 9, 9, 9] |
| 145 | |
| 146 | # Amortize a 5% loan of 1000 with 4 annual payments of 90 |
| 147 | >>> cashflows = [1000, -90, -90, -90, -90] |
| 148 | >>> list(accumulate(cashflows, lambda bal, pmt: bal*1.05 + pmt)) |
| 149 | [1000, 960.0, 918.0, 873.9000000000001, 827.5950000000001] |
| 150 | |
Georg Brandl | 5d94134 | 2016-02-26 19:37:12 +0100 | [diff] [blame] | 151 | # Chaotic recurrence relation https://en.wikipedia.org/wiki/Logistic_map |
Raymond Hettinger | 295c1d4 | 2011-04-21 11:09:28 -0700 | [diff] [blame] | 152 | >>> logistic_map = lambda x, _: r * x * (1 - x) |
| 153 | >>> r = 3.8 |
| 154 | >>> x0 = 0.4 |
| 155 | >>> inputs = repeat(x0, 36) # only the initial value is used |
| 156 | >>> [format(x, '.2f') for x in accumulate(inputs, logistic_map)] |
| 157 | ['0.40', '0.91', '0.30', '0.81', '0.60', '0.92', '0.29', '0.79', '0.63', |
Serhiy Storchaka | a4d170d | 2013-12-23 18:20:51 +0200 | [diff] [blame] | 158 | '0.88', '0.39', '0.90', '0.33', '0.84', '0.52', '0.95', '0.18', '0.57', |
Raymond Hettinger | 295c1d4 | 2011-04-21 11:09:28 -0700 | [diff] [blame] | 159 | '0.93', '0.25', '0.71', '0.79', '0.63', '0.88', '0.39', '0.91', '0.32', |
| 160 | '0.83', '0.54', '0.95', '0.20', '0.60', '0.91', '0.30', '0.80', '0.60'] |
| 161 | |
Raymond Hettinger | 6480168 | 2013-10-12 16:04:17 -0700 | [diff] [blame] | 162 | See :func:`functools.reduce` for a similar function that returns only the |
| 163 | final accumulated value. |
| 164 | |
Raymond Hettinger | adb8146 | 2010-12-01 22:50:36 +0000 | [diff] [blame] | 165 | .. versionadded:: 3.2 |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 166 | |
Raymond Hettinger | 5d44613 | 2011-03-27 18:52:10 -0700 | [diff] [blame] | 167 | .. versionchanged:: 3.3 |
| 168 | Added the optional *func* parameter. |
| 169 | |
Lisa Roach | 9718b59 | 2018-09-23 17:34:59 -0700 | [diff] [blame] | 170 | .. versionchanged:: 3.8 |
| 171 | Added the optional *initial* parameter. |
| 172 | |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 173 | .. function:: chain(*iterables) |
| 174 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 175 | Make an iterator that returns elements from the first iterable until it is |
| 176 | exhausted, then proceeds to the next iterable, until all of the iterables are |
| 177 | exhausted. Used for treating consecutive sequences as a single sequence. |
Raymond Hettinger | 672866d | 2016-05-28 00:17:54 -0700 | [diff] [blame] | 178 | Roughly equivalent to:: |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 179 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 180 | def chain(*iterables): |
| 181 | # chain('ABC', 'DEF') --> A B C D E F |
| 182 | for it in iterables: |
| 183 | for element in it: |
| 184 | yield element |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 185 | |
| 186 | |
Georg Brandl | 933b974 | 2010-07-29 14:36:11 +0000 | [diff] [blame] | 187 | .. classmethod:: chain.from_iterable(iterable) |
Christian Heimes | 70e7ea2 | 2008-02-28 20:02:27 +0000 | [diff] [blame] | 188 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 189 | Alternate constructor for :func:`chain`. Gets chained inputs from a |
Raymond Hettinger | 1e21ebc | 2013-09-09 01:54:27 -0500 | [diff] [blame] | 190 | single iterable argument that is evaluated lazily. Roughly equivalent to:: |
Christian Heimes | 70e7ea2 | 2008-02-28 20:02:27 +0000 | [diff] [blame] | 191 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 192 | def from_iterable(iterables): |
| 193 | # chain.from_iterable(['ABC', 'DEF']) --> A B C D E F |
| 194 | for it in iterables: |
| 195 | for element in it: |
| 196 | yield element |
Christian Heimes | 70e7ea2 | 2008-02-28 20:02:27 +0000 | [diff] [blame] | 197 | |
Christian Heimes | 7864476 | 2008-03-04 23:39:23 +0000 | [diff] [blame] | 198 | |
Christian Heimes | 836baa5 | 2008-02-26 08:18:30 +0000 | [diff] [blame] | 199 | .. function:: combinations(iterable, r) |
| 200 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 201 | Return *r* length subsequences of elements from the input *iterable*. |
Christian Heimes | 836baa5 | 2008-02-26 08:18:30 +0000 | [diff] [blame] | 202 | |
Ruaridh Williamson | 5e0ed8a | 2020-05-28 20:56:43 +0100 | [diff] [blame] | 203 | The combination tuples are emitted in lexicographic ordering according to |
| 204 | the order of the input *iterable*. So, if the input *iterable* is sorted, |
| 205 | the combination tuples will be produced in sorted order. |
Christian Heimes | 836baa5 | 2008-02-26 08:18:30 +0000 | [diff] [blame] | 206 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 207 | Elements are treated as unique based on their position, not on their |
| 208 | value. So if the input elements are unique, there will be no repeat |
| 209 | values in each combination. |
Christian Heimes | 836baa5 | 2008-02-26 08:18:30 +0000 | [diff] [blame] | 210 | |
Raymond Hettinger | 672866d | 2016-05-28 00:17:54 -0700 | [diff] [blame] | 211 | Roughly equivalent to:: |
Christian Heimes | 836baa5 | 2008-02-26 08:18:30 +0000 | [diff] [blame] | 212 | |
| 213 | def combinations(iterable, r): |
Raymond Hettinger | dd1150e | 2008-03-13 02:39:40 +0000 | [diff] [blame] | 214 | # combinations('ABCD', 2) --> AB AC AD BC BD CD |
| 215 | # combinations(range(4), 3) --> 012 013 023 123 |
Christian Heimes | 836baa5 | 2008-02-26 08:18:30 +0000 | [diff] [blame] | 216 | pool = tuple(iterable) |
Christian Heimes | 380f7f2 | 2008-02-28 11:19:05 +0000 | [diff] [blame] | 217 | n = len(pool) |
Raymond Hettinger | 5bad41e | 2009-01-08 21:01:54 +0000 | [diff] [blame] | 218 | if r > n: |
| 219 | return |
| 220 | indices = list(range(r)) |
Christian Heimes | b558a2e | 2008-03-02 22:46:37 +0000 | [diff] [blame] | 221 | yield tuple(pool[i] for i in indices) |
Raymond Hettinger | cf984ce | 2009-02-18 20:56:51 +0000 | [diff] [blame] | 222 | while True: |
Christian Heimes | 380f7f2 | 2008-02-28 11:19:05 +0000 | [diff] [blame] | 223 | for i in reversed(range(r)): |
Christian Heimes | b558a2e | 2008-03-02 22:46:37 +0000 | [diff] [blame] | 224 | if indices[i] != i + n - r: |
Christian Heimes | 836baa5 | 2008-02-26 08:18:30 +0000 | [diff] [blame] | 225 | break |
Christian Heimes | 380f7f2 | 2008-02-28 11:19:05 +0000 | [diff] [blame] | 226 | else: |
| 227 | return |
Christian Heimes | b558a2e | 2008-03-02 22:46:37 +0000 | [diff] [blame] | 228 | indices[i] += 1 |
Christian Heimes | 380f7f2 | 2008-02-28 11:19:05 +0000 | [diff] [blame] | 229 | for j in range(i+1, r): |
Christian Heimes | b558a2e | 2008-03-02 22:46:37 +0000 | [diff] [blame] | 230 | indices[j] = indices[j-1] + 1 |
| 231 | yield tuple(pool[i] for i in indices) |
Christian Heimes | 836baa5 | 2008-02-26 08:18:30 +0000 | [diff] [blame] | 232 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 233 | The code for :func:`combinations` can be also expressed as a subsequence |
| 234 | of :func:`permutations` after filtering entries where the elements are not |
| 235 | in sorted order (according to their position in the input pool):: |
Christian Heimes | 7864476 | 2008-03-04 23:39:23 +0000 | [diff] [blame] | 236 | |
| 237 | def combinations(iterable, r): |
| 238 | pool = tuple(iterable) |
| 239 | n = len(pool) |
| 240 | for indices in permutations(range(n), r): |
| 241 | if sorted(indices) == list(indices): |
| 242 | yield tuple(pool[i] for i in indices) |
| 243 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 244 | The number of items returned is ``n! / r! / (n-r)!`` when ``0 <= r <= n`` |
| 245 | or zero when ``r > n``. |
Christian Heimes | 836baa5 | 2008-02-26 08:18:30 +0000 | [diff] [blame] | 246 | |
Raymond Hettinger | d07d939 | 2009-01-27 04:20:44 +0000 | [diff] [blame] | 247 | .. function:: combinations_with_replacement(iterable, r) |
| 248 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 249 | Return *r* length subsequences of elements from the input *iterable* |
| 250 | allowing individual elements to be repeated more than once. |
Raymond Hettinger | d07d939 | 2009-01-27 04:20:44 +0000 | [diff] [blame] | 251 | |
Ruaridh Williamson | 5e0ed8a | 2020-05-28 20:56:43 +0100 | [diff] [blame] | 252 | The combination tuples are emitted in lexicographic ordering according to |
| 253 | the order of the input *iterable*. So, if the input *iterable* is sorted, |
| 254 | the combination tuples will be produced in sorted order. |
Raymond Hettinger | d07d939 | 2009-01-27 04:20:44 +0000 | [diff] [blame] | 255 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 256 | Elements are treated as unique based on their position, not on their |
| 257 | value. So if the input elements are unique, the generated combinations |
| 258 | will also be unique. |
Raymond Hettinger | d07d939 | 2009-01-27 04:20:44 +0000 | [diff] [blame] | 259 | |
Raymond Hettinger | 672866d | 2016-05-28 00:17:54 -0700 | [diff] [blame] | 260 | Roughly equivalent to:: |
Raymond Hettinger | d07d939 | 2009-01-27 04:20:44 +0000 | [diff] [blame] | 261 | |
| 262 | def combinations_with_replacement(iterable, r): |
| 263 | # combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC |
| 264 | pool = tuple(iterable) |
| 265 | n = len(pool) |
| 266 | if not n and r: |
| 267 | return |
| 268 | indices = [0] * r |
| 269 | yield tuple(pool[i] for i in indices) |
Raymond Hettinger | cf984ce | 2009-02-18 20:56:51 +0000 | [diff] [blame] | 270 | while True: |
Raymond Hettinger | d07d939 | 2009-01-27 04:20:44 +0000 | [diff] [blame] | 271 | for i in reversed(range(r)): |
| 272 | if indices[i] != n - 1: |
| 273 | break |
| 274 | else: |
| 275 | return |
| 276 | indices[i:] = [indices[i] + 1] * (r - i) |
| 277 | yield tuple(pool[i] for i in indices) |
| 278 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 279 | The code for :func:`combinations_with_replacement` can be also expressed as |
| 280 | a subsequence of :func:`product` after filtering entries where the elements |
| 281 | are not in sorted order (according to their position in the input pool):: |
Raymond Hettinger | d07d939 | 2009-01-27 04:20:44 +0000 | [diff] [blame] | 282 | |
| 283 | def combinations_with_replacement(iterable, r): |
| 284 | pool = tuple(iterable) |
| 285 | n = len(pool) |
| 286 | for indices in product(range(n), repeat=r): |
| 287 | if sorted(indices) == list(indices): |
| 288 | yield tuple(pool[i] for i in indices) |
| 289 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 290 | The number of items returned is ``(n+r-1)! / r! / (n-1)!`` when ``n > 0``. |
Raymond Hettinger | d07d939 | 2009-01-27 04:20:44 +0000 | [diff] [blame] | 291 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 292 | .. versionadded:: 3.1 |
Raymond Hettinger | d07d939 | 2009-01-27 04:20:44 +0000 | [diff] [blame] | 293 | |
Georg Brandl | 67b21b7 | 2010-08-17 15:07:14 +0000 | [diff] [blame] | 294 | |
Raymond Hettinger | 6b3b0fc | 2009-01-26 02:56:58 +0000 | [diff] [blame] | 295 | .. function:: compress(data, selectors) |
| 296 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 297 | Make an iterator that filters elements from *data* returning only those that |
| 298 | have a corresponding element in *selectors* that evaluates to ``True``. |
| 299 | Stops when either the *data* or *selectors* iterables has been exhausted. |
Raymond Hettinger | 672866d | 2016-05-28 00:17:54 -0700 | [diff] [blame] | 300 | Roughly equivalent to:: |
Raymond Hettinger | 6b3b0fc | 2009-01-26 02:56:58 +0000 | [diff] [blame] | 301 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 302 | def compress(data, selectors): |
| 303 | # compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F |
| 304 | return (d for d, s in zip(data, selectors) if s) |
Raymond Hettinger | 6b3b0fc | 2009-01-26 02:56:58 +0000 | [diff] [blame] | 305 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 306 | .. versionadded:: 3.1 |
Raymond Hettinger | 6b3b0fc | 2009-01-26 02:56:58 +0000 | [diff] [blame] | 307 | |
| 308 | |
Raymond Hettinger | 9e8dbbc | 2009-02-14 04:21:49 +0000 | [diff] [blame] | 309 | .. function:: count(start=0, step=1) |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 310 | |
Andrew Kuchling | edb4260 | 2013-06-21 08:00:58 -0400 | [diff] [blame] | 311 | Make an iterator that returns evenly spaced values starting with number *start*. Often |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 312 | used as an argument to :func:`map` to generate consecutive data points. |
Raymond Hettinger | 672866d | 2016-05-28 00:17:54 -0700 | [diff] [blame] | 313 | Also, used with :func:`zip` to add sequence numbers. Roughly equivalent to:: |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 314 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 315 | def count(start=0, step=1): |
| 316 | # count(10) --> 10 11 12 13 14 ... |
Georg Brandl | 37a80dc | 2011-01-13 07:31:18 +0000 | [diff] [blame] | 317 | # count(2.5, 0.5) -> 2.5 3.0 3.5 ... |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 318 | n = start |
| 319 | while True: |
| 320 | yield n |
| 321 | n += step |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 322 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 323 | When counting with floating point numbers, better accuracy can sometimes be |
| 324 | achieved by substituting multiplicative code such as: ``(start + step * i |
| 325 | for i in count())``. |
Raymond Hettinger | 5bc472a | 2009-06-17 01:40:52 +0000 | [diff] [blame] | 326 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 327 | .. versionchanged:: 3.1 |
| 328 | Added *step* argument and allowed non-integer arguments. |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 329 | |
| 330 | .. function:: cycle(iterable) |
| 331 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 332 | Make an iterator returning elements from the iterable and saving a copy of each. |
| 333 | When the iterable is exhausted, return elements from the saved copy. Repeats |
Raymond Hettinger | 672866d | 2016-05-28 00:17:54 -0700 | [diff] [blame] | 334 | indefinitely. Roughly equivalent to:: |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 335 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 336 | def cycle(iterable): |
| 337 | # cycle('ABCD') --> A B C D A B C D A B C D ... |
| 338 | saved = [] |
| 339 | for element in iterable: |
| 340 | yield element |
| 341 | saved.append(element) |
| 342 | while saved: |
| 343 | for element in saved: |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 344 | yield element |
| 345 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 346 | Note, this member of the toolkit may require significant auxiliary storage |
| 347 | (depending on the length of the iterable). |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 348 | |
| 349 | |
| 350 | .. function:: dropwhile(predicate, iterable) |
| 351 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 352 | Make an iterator that drops elements from the iterable as long as the predicate |
| 353 | is true; afterwards, returns every element. Note, the iterator does not produce |
| 354 | *any* output until the predicate first becomes false, so it may have a lengthy |
Raymond Hettinger | 672866d | 2016-05-28 00:17:54 -0700 | [diff] [blame] | 355 | start-up time. Roughly equivalent to:: |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 356 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 357 | def dropwhile(predicate, iterable): |
| 358 | # dropwhile(lambda x: x<5, [1,4,6,4,1]) --> 6 4 1 |
| 359 | iterable = iter(iterable) |
| 360 | for x in iterable: |
| 361 | if not predicate(x): |
| 362 | yield x |
| 363 | break |
| 364 | for x in iterable: |
| 365 | yield x |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 366 | |
Raymond Hettinger | 749761e | 2009-01-27 04:42:48 +0000 | [diff] [blame] | 367 | .. function:: filterfalse(predicate, iterable) |
| 368 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 369 | Make an iterator that filters elements from iterable returning only those for |
| 370 | which the predicate is ``False``. If *predicate* is ``None``, return the items |
Raymond Hettinger | 672866d | 2016-05-28 00:17:54 -0700 | [diff] [blame] | 371 | that are false. Roughly equivalent to:: |
Raymond Hettinger | 749761e | 2009-01-27 04:42:48 +0000 | [diff] [blame] | 372 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 373 | def filterfalse(predicate, iterable): |
| 374 | # filterfalse(lambda x: x%2, range(10)) --> 0 2 4 6 8 |
| 375 | if predicate is None: |
| 376 | predicate = bool |
| 377 | for x in iterable: |
| 378 | if not predicate(x): |
| 379 | yield x |
Raymond Hettinger | 749761e | 2009-01-27 04:42:48 +0000 | [diff] [blame] | 380 | |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 381 | |
Georg Brandl | 3dd3388 | 2009-06-01 17:35:27 +0000 | [diff] [blame] | 382 | .. function:: groupby(iterable, key=None) |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 383 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 384 | Make an iterator that returns consecutive keys and groups from the *iterable*. |
| 385 | The *key* is a function computing a key value for each element. If not |
| 386 | specified or is ``None``, *key* defaults to an identity function and returns |
| 387 | the element unchanged. Generally, the iterable needs to already be sorted on |
| 388 | the same key function. |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 389 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 390 | The operation of :func:`groupby` is similar to the ``uniq`` filter in Unix. It |
| 391 | generates a break or new group every time the value of the key function changes |
| 392 | (which is why it is usually necessary to have sorted the data using the same key |
| 393 | function). That behavior differs from SQL's GROUP BY which aggregates common |
| 394 | elements regardless of their input order. |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 395 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 396 | The returned group is itself an iterator that shares the underlying iterable |
| 397 | with :func:`groupby`. Because the source is shared, when the :func:`groupby` |
| 398 | object is advanced, the previous group is no longer visible. So, if that data |
| 399 | is needed later, it should be stored as a list:: |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 400 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 401 | groups = [] |
| 402 | uniquekeys = [] |
| 403 | data = sorted(data, key=keyfunc) |
| 404 | for k, g in groupby(data, keyfunc): |
| 405 | groups.append(list(g)) # Store group iterator as a list |
| 406 | uniquekeys.append(k) |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 407 | |
Raymond Hettinger | 672866d | 2016-05-28 00:17:54 -0700 | [diff] [blame] | 408 | :func:`groupby` is roughly equivalent to:: |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 409 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 410 | class groupby: |
| 411 | # [k for k, g in groupby('AAAABBBCCDAABBB')] --> A B C D A B |
| 412 | # [list(g) for k, g in groupby('AAAABBBCCD')] --> AAAA BBB CC D |
| 413 | def __init__(self, iterable, key=None): |
| 414 | if key is None: |
| 415 | key = lambda x: x |
| 416 | self.keyfunc = key |
| 417 | self.it = iter(iterable) |
| 418 | self.tgtkey = self.currkey = self.currvalue = object() |
| 419 | def __iter__(self): |
| 420 | return self |
| 421 | def __next__(self): |
Serhiy Storchaka | c247caf | 2017-09-24 13:36:11 +0300 | [diff] [blame] | 422 | self.id = object() |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 423 | while self.currkey == self.tgtkey: |
| 424 | self.currvalue = next(self.it) # Exit on StopIteration |
| 425 | self.currkey = self.keyfunc(self.currvalue) |
| 426 | self.tgtkey = self.currkey |
Serhiy Storchaka | c247caf | 2017-09-24 13:36:11 +0300 | [diff] [blame] | 427 | return (self.currkey, self._grouper(self.tgtkey, self.id)) |
| 428 | def _grouper(self, tgtkey, id): |
| 429 | while self.id is id and self.currkey == tgtkey: |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 430 | yield self.currvalue |
Raymond Hettinger | 828d932 | 2014-11-22 21:56:23 -0800 | [diff] [blame] | 431 | try: |
| 432 | self.currvalue = next(self.it) |
| 433 | except StopIteration: |
| 434 | return |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 435 | self.currkey = self.keyfunc(self.currvalue) |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 436 | |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 437 | |
Ezio Melotti | e0add76 | 2012-09-14 06:32:35 +0300 | [diff] [blame] | 438 | .. function:: islice(iterable, stop) |
| 439 | islice(iterable, start, stop[, step]) |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 440 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 441 | Make an iterator that returns selected elements from the iterable. If *start* is |
| 442 | non-zero, then elements from the iterable are skipped until start is reached. |
| 443 | Afterward, elements are returned consecutively unless *step* is set higher than |
| 444 | one which results in items being skipped. If *stop* is ``None``, then iteration |
| 445 | continues until the iterator is exhausted, if at all; otherwise, it stops at the |
| 446 | specified position. Unlike regular slicing, :func:`islice` does not support |
| 447 | negative values for *start*, *stop*, or *step*. Can be used to extract related |
| 448 | fields from data where the internal structure has been flattened (for example, a |
Raymond Hettinger | 672866d | 2016-05-28 00:17:54 -0700 | [diff] [blame] | 449 | multi-line report may list a name field on every third line). Roughly equivalent to:: |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 450 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 451 | def islice(iterable, *args): |
| 452 | # islice('ABCDEFG', 2) --> A B |
| 453 | # islice('ABCDEFG', 2, 4) --> C D |
| 454 | # islice('ABCDEFG', 2, None) --> C D E F G |
| 455 | # islice('ABCDEFG', 0, None, 2) --> A C E G |
| 456 | s = slice(*args) |
Cheryl Sabella | da1734c | 2018-03-26 21:29:33 -0400 | [diff] [blame] | 457 | start, stop, step = s.start or 0, s.stop or sys.maxsize, s.step or 1 |
| 458 | it = iter(range(start, stop, step)) |
Raymond Hettinger | 828d932 | 2014-11-22 21:56:23 -0800 | [diff] [blame] | 459 | try: |
| 460 | nexti = next(it) |
| 461 | except StopIteration: |
Cheryl Sabella | da1734c | 2018-03-26 21:29:33 -0400 | [diff] [blame] | 462 | # Consume *iterable* up to the *start* position. |
| 463 | for i, element in zip(range(start), iterable): |
| 464 | pass |
Raymond Hettinger | 828d932 | 2014-11-22 21:56:23 -0800 | [diff] [blame] | 465 | return |
Cheryl Sabella | da1734c | 2018-03-26 21:29:33 -0400 | [diff] [blame] | 466 | try: |
| 467 | for i, element in enumerate(iterable): |
| 468 | if i == nexti: |
| 469 | yield element |
| 470 | nexti = next(it) |
| 471 | except StopIteration: |
| 472 | # Consume to *stop*. |
| 473 | for i, element in zip(range(i + 1, stop), iterable): |
| 474 | pass |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 475 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 476 | If *start* is ``None``, then iteration starts at zero. If *step* is ``None``, |
| 477 | then the step defaults to one. |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 478 | |
Raymond Hettinger | cc061d0 | 2020-11-30 20:42:54 -0800 | [diff] [blame] | 479 | .. function:: pairwise(iterable) |
| 480 | |
| 481 | Return successive overlapping pairs taken from the input *iterable*. |
| 482 | |
| 483 | The number of 2-tuples in the output iterator will be one fewer than the |
| 484 | number of inputs. It will be empty if the input iterable has fewer than |
| 485 | two values. |
| 486 | |
| 487 | Roughly equivalent to:: |
| 488 | |
| 489 | def pairwise(iterable): |
| 490 | # pairwise('ABCDEFG') --> AB BC CD DE EF FG |
| 491 | a, b = tee(iterable) |
| 492 | next(b, None) |
| 493 | return zip(a, b) |
| 494 | |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 495 | |
Georg Brandl | 3dd3388 | 2009-06-01 17:35:27 +0000 | [diff] [blame] | 496 | .. function:: permutations(iterable, r=None) |
Christian Heimes | 70e7ea2 | 2008-02-28 20:02:27 +0000 | [diff] [blame] | 497 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 498 | Return successive *r* length permutations of elements in the *iterable*. |
Christian Heimes | 70e7ea2 | 2008-02-28 20:02:27 +0000 | [diff] [blame] | 499 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 500 | If *r* is not specified or is ``None``, then *r* defaults to the length |
| 501 | of the *iterable* and all possible full-length permutations |
| 502 | are generated. |
Christian Heimes | 70e7ea2 | 2008-02-28 20:02:27 +0000 | [diff] [blame] | 503 | |
Ruaridh Williamson | 5e0ed8a | 2020-05-28 20:56:43 +0100 | [diff] [blame] | 504 | The permutation tuples are emitted in lexicographic ordering according to |
| 505 | the order of the input *iterable*. So, if the input *iterable* is sorted, |
| 506 | the combination tuples will be produced in sorted order. |
Christian Heimes | 70e7ea2 | 2008-02-28 20:02:27 +0000 | [diff] [blame] | 507 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 508 | Elements are treated as unique based on their position, not on their |
| 509 | value. So if the input elements are unique, there will be no repeat |
| 510 | values in each permutation. |
Christian Heimes | 70e7ea2 | 2008-02-28 20:02:27 +0000 | [diff] [blame] | 511 | |
Raymond Hettinger | 672866d | 2016-05-28 00:17:54 -0700 | [diff] [blame] | 512 | Roughly equivalent to:: |
Christian Heimes | b558a2e | 2008-03-02 22:46:37 +0000 | [diff] [blame] | 513 | |
| 514 | def permutations(iterable, r=None): |
Raymond Hettinger | dd1150e | 2008-03-13 02:39:40 +0000 | [diff] [blame] | 515 | # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC |
| 516 | # permutations(range(3)) --> 012 021 102 120 201 210 |
Christian Heimes | b558a2e | 2008-03-02 22:46:37 +0000 | [diff] [blame] | 517 | pool = tuple(iterable) |
| 518 | n = len(pool) |
| 519 | r = n if r is None else r |
Raymond Hettinger | 5bad41e | 2009-01-08 21:01:54 +0000 | [diff] [blame] | 520 | if r > n: |
| 521 | return |
| 522 | indices = list(range(n)) |
Sandro Tosi | 7386662 | 2011-12-25 17:25:45 +0100 | [diff] [blame] | 523 | cycles = list(range(n, n-r, -1)) |
Christian Heimes | b558a2e | 2008-03-02 22:46:37 +0000 | [diff] [blame] | 524 | yield tuple(pool[i] for i in indices[:r]) |
| 525 | while n: |
| 526 | for i in reversed(range(r)): |
| 527 | cycles[i] -= 1 |
| 528 | if cycles[i] == 0: |
| 529 | indices[i:] = indices[i+1:] + indices[i:i+1] |
| 530 | cycles[i] = n - i |
| 531 | else: |
| 532 | j = cycles[i] |
| 533 | indices[i], indices[-j] = indices[-j], indices[i] |
| 534 | yield tuple(pool[i] for i in indices[:r]) |
| 535 | break |
| 536 | else: |
| 537 | return |
Christian Heimes | 70e7ea2 | 2008-02-28 20:02:27 +0000 | [diff] [blame] | 538 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 539 | The code for :func:`permutations` can be also expressed as a subsequence of |
| 540 | :func:`product`, filtered to exclude entries with repeated elements (those |
| 541 | from the same position in the input pool):: |
Christian Heimes | 7864476 | 2008-03-04 23:39:23 +0000 | [diff] [blame] | 542 | |
| 543 | def permutations(iterable, r=None): |
| 544 | pool = tuple(iterable) |
| 545 | n = len(pool) |
| 546 | r = n if r is None else r |
| 547 | for indices in product(range(n), repeat=r): |
| 548 | if len(set(indices)) == r: |
| 549 | yield tuple(pool[i] for i in indices) |
| 550 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 551 | The number of items returned is ``n! / (n-r)!`` when ``0 <= r <= n`` |
| 552 | or zero when ``r > n``. |
Christian Heimes | 70e7ea2 | 2008-02-28 20:02:27 +0000 | [diff] [blame] | 553 | |
Georg Brandl | 3dd3388 | 2009-06-01 17:35:27 +0000 | [diff] [blame] | 554 | .. function:: product(*iterables, repeat=1) |
Christian Heimes | 90c3d9b | 2008-02-23 13:18:03 +0000 | [diff] [blame] | 555 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 556 | Cartesian product of input iterables. |
Christian Heimes | 90c3d9b | 2008-02-23 13:18:03 +0000 | [diff] [blame] | 557 | |
Raymond Hettinger | 672866d | 2016-05-28 00:17:54 -0700 | [diff] [blame] | 558 | Roughly equivalent to nested for-loops in a generator expression. For example, |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 559 | ``product(A, B)`` returns the same as ``((x,y) for x in A for y in B)``. |
Christian Heimes | 90c3d9b | 2008-02-23 13:18:03 +0000 | [diff] [blame] | 560 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 561 | The nested loops cycle like an odometer with the rightmost element advancing |
| 562 | on every iteration. This pattern creates a lexicographic ordering so that if |
| 563 | the input's iterables are sorted, the product tuples are emitted in sorted |
| 564 | order. |
Christian Heimes | 90c3d9b | 2008-02-23 13:18:03 +0000 | [diff] [blame] | 565 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 566 | To compute the product of an iterable with itself, specify the number of |
| 567 | repetitions with the optional *repeat* keyword argument. For example, |
| 568 | ``product(A, repeat=4)`` means the same as ``product(A, A, A, A)``. |
Christian Heimes | 9e7f1d2 | 2008-02-28 12:27:11 +0000 | [diff] [blame] | 569 | |
Raymond Hettinger | 672866d | 2016-05-28 00:17:54 -0700 | [diff] [blame] | 570 | This function is roughly equivalent to the following code, except that the |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 571 | actual implementation does not build up intermediate results in memory:: |
Christian Heimes | 90c3d9b | 2008-02-23 13:18:03 +0000 | [diff] [blame] | 572 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 573 | def product(*args, repeat=1): |
| 574 | # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy |
| 575 | # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111 |
| 576 | pools = [tuple(pool) for pool in args] * repeat |
| 577 | result = [[]] |
| 578 | for pool in pools: |
| 579 | result = [x+[y] for x in result for y in pool] |
| 580 | for prod in result: |
| 581 | yield tuple(prod) |
Christian Heimes | 90c3d9b | 2008-02-23 13:18:03 +0000 | [diff] [blame] | 582 | |
Ramil Nugmanov | cfc6ce4 | 2020-05-28 19:46:22 +0300 | [diff] [blame] | 583 | Before :func:`product` runs, it completely consumes the input iterables, |
| 584 | keeping pools of values in memory to generate the products. Accordingly, |
Zackery Spytz | 074ad51 | 2020-12-17 13:24:43 -0700 | [diff] [blame] | 585 | it is only useful with finite inputs. |
Christian Heimes | 90c3d9b | 2008-02-23 13:18:03 +0000 | [diff] [blame] | 586 | |
Raymond Hettinger | d75ad44 | 2009-06-01 19:16:52 +0000 | [diff] [blame] | 587 | .. function:: repeat(object[, times]) |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 588 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 589 | Make an iterator that returns *object* over and over again. Runs indefinitely |
| 590 | unless the *times* argument is specified. Used as argument to :func:`map` for |
| 591 | invariant parameters to the called function. Also used with :func:`zip` to |
Raymond Hettinger | 672866d | 2016-05-28 00:17:54 -0700 | [diff] [blame] | 592 | create an invariant part of a tuple record. |
| 593 | |
| 594 | Roughly equivalent to:: |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 595 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 596 | def repeat(object, times=None): |
| 597 | # repeat(10, 3) --> 10 10 10 |
| 598 | if times is None: |
| 599 | while True: |
| 600 | yield object |
| 601 | else: |
| 602 | for i in range(times): |
| 603 | yield object |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 604 | |
Raymond Hettinger | fc3ba6b | 2012-02-01 09:07:40 -0800 | [diff] [blame] | 605 | A common use for *repeat* is to supply a stream of constant values to *map* |
| 606 | or *zip*:: |
| 607 | |
| 608 | >>> list(map(pow, range(10), repeat(2))) |
| 609 | [0, 1, 4, 9, 16, 25, 36, 49, 64, 81] |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 610 | |
| 611 | .. function:: starmap(function, iterable) |
| 612 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 613 | Make an iterator that computes the function using arguments obtained from |
| 614 | the iterable. Used instead of :func:`map` when argument parameters are already |
| 615 | grouped in tuples from a single iterable (the data has been "pre-zipped"). The |
| 616 | difference between :func:`map` and :func:`starmap` parallels the distinction |
Raymond Hettinger | 672866d | 2016-05-28 00:17:54 -0700 | [diff] [blame] | 617 | between ``function(a,b)`` and ``function(*c)``. Roughly equivalent to:: |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 618 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 619 | def starmap(function, iterable): |
| 620 | # starmap(pow, [(2,5), (3,2), (10,3)]) --> 32 9 1000 |
| 621 | for args in iterable: |
| 622 | yield function(*args) |
Christian Heimes | 679db4a | 2008-01-18 09:56:22 +0000 | [diff] [blame] | 623 | |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 624 | |
| 625 | .. function:: takewhile(predicate, iterable) |
| 626 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 627 | Make an iterator that returns elements from the iterable as long as the |
Raymond Hettinger | 672866d | 2016-05-28 00:17:54 -0700 | [diff] [blame] | 628 | predicate is true. Roughly equivalent to:: |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 629 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 630 | def takewhile(predicate, iterable): |
| 631 | # takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4 |
| 632 | for x in iterable: |
| 633 | if predicate(x): |
| 634 | yield x |
| 635 | else: |
| 636 | break |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 637 | |
| 638 | |
Georg Brandl | 3dd3388 | 2009-06-01 17:35:27 +0000 | [diff] [blame] | 639 | .. function:: tee(iterable, n=2) |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 640 | |
Raymond Hettinger | ab425aa | 2016-04-26 00:10:00 -0700 | [diff] [blame] | 641 | Return *n* independent iterators from a single iterable. |
| 642 | |
| 643 | The following Python code helps explain what *tee* does (although the actual |
Serhiy Storchaka | 4ecfa45 | 2016-05-16 09:31:54 +0300 | [diff] [blame] | 644 | implementation is more complex and uses only a single underlying |
Raymond Hettinger | 672866d | 2016-05-28 00:17:54 -0700 | [diff] [blame] | 645 | :abbr:`FIFO (first-in, first-out)` queue). |
| 646 | |
| 647 | Roughly equivalent to:: |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 648 | |
Raymond Hettinger | cf984ce | 2009-02-18 20:56:51 +0000 | [diff] [blame] | 649 | def tee(iterable, n=2): |
| 650 | it = iter(iterable) |
| 651 | deques = [collections.deque() for i in range(n)] |
| 652 | def gen(mydeque): |
| 653 | while True: |
| 654 | if not mydeque: # when the local deque is empty |
Raymond Hettinger | 828d932 | 2014-11-22 21:56:23 -0800 | [diff] [blame] | 655 | try: |
| 656 | newval = next(it) # fetch a new value and |
| 657 | except StopIteration: |
| 658 | return |
Raymond Hettinger | cf984ce | 2009-02-18 20:56:51 +0000 | [diff] [blame] | 659 | for d in deques: # load it to all the deques |
| 660 | d.append(newval) |
| 661 | yield mydeque.popleft() |
| 662 | return tuple(gen(d) for d in deques) |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 663 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 664 | Once :func:`tee` has made a split, the original *iterable* should not be |
| 665 | used anywhere else; otherwise, the *iterable* could get advanced without |
Serhiy Storchaka | 918b468 | 2019-09-09 11:18:16 +0300 | [diff] [blame] | 666 | the tee objects being informed. |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 667 | |
Serhiy Storchaka | 526a014 | 2019-09-09 11:47:14 +0300 | [diff] [blame] | 668 | ``tee`` iterators are not threadsafe. A :exc:`RuntimeError` may be |
| 669 | raised when using simultaneously iterators returned by the same :func:`tee` |
| 670 | call, even if the original *iterable* is threadsafe. |
| 671 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 672 | This itertool may require significant auxiliary storage (depending on how |
| 673 | much temporary data needs to be stored). In general, if one iterator uses |
| 674 | most or all of the data before another iterator starts, it is faster to use |
| 675 | :func:`list` instead of :func:`tee`. |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 676 | |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 677 | |
Georg Brandl | 3dd3388 | 2009-06-01 17:35:27 +0000 | [diff] [blame] | 678 | .. function:: zip_longest(*iterables, fillvalue=None) |
Raymond Hettinger | 749761e | 2009-01-27 04:42:48 +0000 | [diff] [blame] | 679 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 680 | Make an iterator that aggregates elements from each of the iterables. If the |
| 681 | iterables are of uneven length, missing values are filled-in with *fillvalue*. |
Raymond Hettinger | 672866d | 2016-05-28 00:17:54 -0700 | [diff] [blame] | 682 | Iteration continues until the longest iterable is exhausted. Roughly equivalent to:: |
Raymond Hettinger | 749761e | 2009-01-27 04:42:48 +0000 | [diff] [blame] | 683 | |
Raymond Hettinger | 3147b04 | 2017-09-07 14:01:44 -0700 | [diff] [blame] | 684 | def zip_longest(*args, fillvalue=None): |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 685 | # zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D- |
Raymond Hettinger | 3147b04 | 2017-09-07 14:01:44 -0700 | [diff] [blame] | 686 | iterators = [iter(it) for it in args] |
| 687 | num_active = len(iterators) |
| 688 | if not num_active: |
| 689 | return |
| 690 | while True: |
| 691 | values = [] |
| 692 | for i, it in enumerate(iterators): |
| 693 | try: |
| 694 | value = next(it) |
| 695 | except StopIteration: |
| 696 | num_active -= 1 |
| 697 | if not num_active: |
| 698 | return |
| 699 | iterators[i] = repeat(fillvalue) |
| 700 | value = fillvalue |
| 701 | values.append(value) |
| 702 | yield tuple(values) |
Raymond Hettinger | 749761e | 2009-01-27 04:42:48 +0000 | [diff] [blame] | 703 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 704 | If one of the iterables is potentially infinite, then the :func:`zip_longest` |
| 705 | function should be wrapped with something that limits the number of calls |
| 706 | (for example :func:`islice` or :func:`takewhile`). If not specified, |
| 707 | *fillvalue* defaults to ``None``. |
Raymond Hettinger | 749761e | 2009-01-27 04:42:48 +0000 | [diff] [blame] | 708 | |
| 709 | |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 710 | .. _itertools-recipes: |
| 711 | |
Raymond Hettinger | 1fa7682 | 2010-12-06 23:31:36 +0000 | [diff] [blame] | 712 | Itertools Recipes |
| 713 | ----------------- |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 714 | |
| 715 | This section shows recipes for creating an extended toolset using the existing |
| 716 | itertools as building blocks. |
| 717 | |
Raymond Hettinger | adf02b3 | 2019-08-04 13:35:58 -0700 | [diff] [blame] | 718 | Substantially all of these recipes and many, many others can be installed from |
| 719 | the `more-itertools project <https://pypi.org/project/more-itertools/>`_ found |
| 720 | on the Python Package Index:: |
| 721 | |
| 722 | pip install more-itertools |
| 723 | |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 724 | The extended tools offer the same high performance as the underlying toolset. |
| 725 | The superior memory performance is kept by processing elements one at a time |
| 726 | rather than bringing the whole iterable into memory all at once. Code volume is |
| 727 | kept small by linking the tools together in a functional style which helps |
| 728 | eliminate temporary variables. High speed is retained by preferring |
Georg Brandl | 9afde1c | 2007-11-01 20:32:30 +0000 | [diff] [blame] | 729 | "vectorized" building blocks over the use of for-loops and :term:`generator`\s |
Christian Heimes | fe337bf | 2008-03-23 21:54:12 +0000 | [diff] [blame] | 730 | which incur interpreter overhead. |
| 731 | |
| 732 | .. testcode:: |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 733 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 734 | def take(n, iterable): |
| 735 | "Return first n items of the iterable as a list" |
| 736 | return list(islice(iterable, n)) |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 737 | |
Raymond Hettinger | 9265dd7 | 2018-04-08 08:44:20 -0700 | [diff] [blame] | 738 | def prepend(value, iterator): |
| 739 | "Prepend a single value in front of an iterator" |
| 740 | # prepend(1, [2, 3, 4]) -> 1 2 3 4 |
| 741 | return chain([value], iterator) |
| 742 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 743 | def tabulate(function, start=0): |
| 744 | "Return function(0), function(1), ..." |
| 745 | return map(function, count(start)) |
Raymond Hettinger | dd1150e | 2008-03-13 02:39:40 +0000 | [diff] [blame] | 746 | |
Raymond Hettinger | dfe098d | 2014-05-25 22:03:56 -0700 | [diff] [blame] | 747 | def tail(n, iterable): |
| 748 | "Return an iterator over the last n items" |
| 749 | # tail(3, 'ABCDEFG') --> E F G |
| 750 | return iter(collections.deque(iterable, maxlen=n)) |
| 751 | |
Cheryl Sabella | da1734c | 2018-03-26 21:29:33 -0400 | [diff] [blame] | 752 | def consume(iterator, n=None): |
| 753 | "Advance the iterator n-steps ahead. If n is None, consume entirely." |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 754 | # Use functions that consume iterators at C speed. |
| 755 | if n is None: |
| 756 | # feed the entire iterator into a zero-length deque |
| 757 | collections.deque(iterator, maxlen=0) |
| 758 | else: |
| 759 | # advance to the empty slice starting at position n |
| 760 | next(islice(iterator, n, n), None) |
Raymond Hettinger | fa00796 | 2009-03-09 11:55:25 +0000 | [diff] [blame] | 761 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 762 | def nth(iterable, n, default=None): |
| 763 | "Returns the nth item or a default value" |
| 764 | return next(islice(iterable, n, None), default) |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 765 | |
Raymond Hettinger | e525ee3 | 2016-03-06 18:11:38 -0800 | [diff] [blame] | 766 | def all_equal(iterable): |
| 767 | "Returns True if all the elements are equal to each other" |
| 768 | g = groupby(iterable) |
| 769 | return next(g, True) and not next(g, False) |
| 770 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 771 | def quantify(iterable, pred=bool): |
| 772 | "Count how many times the predicate is true" |
| 773 | return sum(map(pred, iterable)) |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 774 | |
Raymond Hettinger | fc40b30 | 2020-11-29 10:47:22 -0800 | [diff] [blame] | 775 | def pad_none(iterable): |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 776 | """Returns the sequence elements and then returns None indefinitely. |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 777 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 778 | Useful for emulating the behavior of the built-in map() function. |
| 779 | """ |
| 780 | return chain(iterable, repeat(None)) |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 781 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 782 | def ncycles(iterable, n): |
| 783 | "Returns the sequence elements n times" |
| 784 | return chain.from_iterable(repeat(tuple(iterable), n)) |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 785 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 786 | def dotproduct(vec1, vec2): |
| 787 | return sum(map(operator.mul, vec1, vec2)) |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 788 | |
Raymond Hettinger | 77fde8d | 2020-12-25 16:43:20 -0800 | [diff] [blame] | 789 | def convolve(signal, kernel): |
| 790 | # See: https://betterexplained.com/articles/intuitive-convolution/ |
| 791 | # convolve(data, [0.25, 0.25, 0.25, 0.25]) --> Moving average (blur) |
| 792 | # convolve(data, [1, -1]) --> 1st finite difference (1st derivative) |
| 793 | # convolve(data, [1, -2, 1]) --> 2nd finite difference (2nd derivative) |
Raymond Hettinger | f421bfc | 2020-12-30 12:51:19 -0800 | [diff] [blame] | 794 | kernel = tuple(kernel)[::-1] |
Raymond Hettinger | 77fde8d | 2020-12-25 16:43:20 -0800 | [diff] [blame] | 795 | n = len(kernel) |
Raymond Hettinger | f421bfc | 2020-12-30 12:51:19 -0800 | [diff] [blame] | 796 | window = collections.deque([0], maxlen=n) * n |
Raymond Hettinger | 77fde8d | 2020-12-25 16:43:20 -0800 | [diff] [blame] | 797 | for x in chain(signal, repeat(0, n-1)): |
| 798 | window.append(x) |
| 799 | yield sum(map(operator.mul, kernel, window)) |
| 800 | |
Борис Верховский | 99b7701 | 2019-11-02 15:09:14 -0400 | [diff] [blame] | 801 | def flatten(list_of_lists): |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 802 | "Flatten one level of nesting" |
Борис Верховский | 99b7701 | 2019-11-02 15:09:14 -0400 | [diff] [blame] | 803 | return chain.from_iterable(list_of_lists) |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 804 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 805 | def repeatfunc(func, times=None, *args): |
| 806 | """Repeat calls to func with specified arguments. |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 807 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 808 | Example: repeatfunc(random.random) |
| 809 | """ |
| 810 | if times is None: |
| 811 | return starmap(func, repeat(args)) |
| 812 | return starmap(func, repeat(args, times)) |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 813 | |
Raymond Hettinger | 44571da | 2013-05-05 19:53:41 -0700 | [diff] [blame] | 814 | def grouper(iterable, n, fillvalue=None): |
Raymond Hettinger | 9ae9473 | 2012-07-08 16:04:03 -0700 | [diff] [blame] | 815 | "Collect data into fixed-length chunks or blocks" |
Raymond Hettinger | 44571da | 2013-05-05 19:53:41 -0700 | [diff] [blame] | 816 | # grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx" |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 817 | args = [iter(iterable)] * n |
| 818 | return zip_longest(*args, fillvalue=fillvalue) |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 819 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 820 | def roundrobin(*iterables): |
| 821 | "roundrobin('ABC', 'D', 'EF') --> A D E B F C" |
| 822 | # Recipe credited to George Sakkis |
Raymond Hettinger | 337cbba | 2017-11-21 00:23:34 -0800 | [diff] [blame] | 823 | num_active = len(iterables) |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 824 | nexts = cycle(iter(it).__next__ for it in iterables) |
Raymond Hettinger | 337cbba | 2017-11-21 00:23:34 -0800 | [diff] [blame] | 825 | while num_active: |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 826 | try: |
| 827 | for next in nexts: |
| 828 | yield next() |
| 829 | except StopIteration: |
Raymond Hettinger | 337cbba | 2017-11-21 00:23:34 -0800 | [diff] [blame] | 830 | # Remove the iterator we just exhausted from the cycle. |
| 831 | num_active -= 1 |
| 832 | nexts = cycle(islice(nexts, num_active)) |
Georg Brandl | 116aa62 | 2007-08-15 14:28:22 +0000 | [diff] [blame] | 833 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 834 | def partition(pred, iterable): |
Raymond Hettinger | fc40b30 | 2020-11-29 10:47:22 -0800 | [diff] [blame] | 835 | "Use a predicate to partition entries into false entries and true entries" |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 836 | # partition(is_odd, range(10)) --> 0 2 4 6 8 and 1 3 5 7 9 |
| 837 | t1, t2 = tee(iterable) |
| 838 | return filterfalse(pred, t1), filter(pred, t2) |
Raymond Hettinger | 5ce0aa2 | 2010-12-01 10:49:19 +0000 | [diff] [blame] | 839 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 840 | def powerset(iterable): |
| 841 | "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)" |
| 842 | s = list(iterable) |
| 843 | return chain.from_iterable(combinations(s, r) for r in range(len(s)+1)) |
Raymond Hettinger | 08d01ee | 2010-08-07 05:36:53 +0000 | [diff] [blame] | 844 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 845 | def unique_everseen(iterable, key=None): |
| 846 | "List unique elements, preserving order. Remember all elements ever seen." |
| 847 | # unique_everseen('AAAABBBCCDAABBB') --> A B C D |
| 848 | # unique_everseen('ABBCcAD', str.lower) --> A B C D |
| 849 | seen = set() |
| 850 | seen_add = seen.add |
| 851 | if key is None: |
| 852 | for element in filterfalse(seen.__contains__, iterable): |
| 853 | seen_add(element) |
| 854 | yield element |
| 855 | else: |
| 856 | for element in iterable: |
| 857 | k = key(element) |
| 858 | if k not in seen: |
| 859 | seen_add(k) |
| 860 | yield element |
Christian Heimes | 90c3d9b | 2008-02-23 13:18:03 +0000 | [diff] [blame] | 861 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 862 | def unique_justseen(iterable, key=None): |
| 863 | "List unique elements, preserving order. Remember only the element just seen." |
| 864 | # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B |
| 865 | # unique_justseen('ABBCcAD', str.lower) --> A B C A D |
Jakub Molinski | b4c7f39 | 2019-04-23 10:30:30 +0200 | [diff] [blame] | 866 | return map(next, map(operator.itemgetter(1), groupby(iterable, key))) |
Raymond Hettinger | ad9d96b | 2009-01-02 21:39:07 +0000 | [diff] [blame] | 867 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 868 | def iter_except(func, exception, first=None): |
| 869 | """ Call a function repeatedly until an exception is raised. |
Raymond Hettinger | fc91aa2 | 2010-03-28 18:27:13 +0000 | [diff] [blame] | 870 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 871 | Converts a call-until-exception interface to an iterator interface. |
Andrew Kuchling | 1d7d580 | 2013-06-20 21:17:41 -0400 | [diff] [blame] | 872 | Like builtins.iter(func, sentinel) but uses an exception instead |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 873 | of a sentinel to end the loop. |
Raymond Hettinger | fc91aa2 | 2010-03-28 18:27:13 +0000 | [diff] [blame] | 874 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 875 | Examples: |
| 876 | iter_except(functools.partial(heappop, h), IndexError) # priority queue iterator |
| 877 | iter_except(d.popitem, KeyError) # non-blocking dict iterator |
| 878 | iter_except(d.popleft, IndexError) # non-blocking deque iterator |
| 879 | iter_except(q.get_nowait, Queue.Empty) # loop over a producer Queue |
| 880 | iter_except(s.pop, KeyError) # non-blocking set iterator |
Raymond Hettinger | fc91aa2 | 2010-03-28 18:27:13 +0000 | [diff] [blame] | 881 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 882 | """ |
| 883 | try: |
| 884 | if first is not None: |
| 885 | yield first() # For database APIs needing an initial cast to db.first() |
Raymond Hettinger | a503f70 | 2016-03-13 00:12:31 -0800 | [diff] [blame] | 886 | while True: |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 887 | yield func() |
| 888 | except exception: |
| 889 | pass |
Raymond Hettinger | fc91aa2 | 2010-03-28 18:27:13 +0000 | [diff] [blame] | 890 | |
Raymond Hettinger | 31b26f6 | 2014-04-02 03:16:42 -0700 | [diff] [blame] | 891 | def first_true(iterable, default=False, pred=None): |
| 892 | """Returns the first true value in the iterable. |
| 893 | |
| 894 | If no true value is found, returns *default* |
| 895 | |
| 896 | If *pred* is not None, returns the first item |
| 897 | for which pred(item) is true. |
| 898 | |
| 899 | """ |
| 900 | # first_true([a,b,c], x) --> a or b or c or x |
| 901 | # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x |
| 902 | return next(filter(pred, iterable), default) |
| 903 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 904 | def random_product(*args, repeat=1): |
| 905 | "Random selection from itertools.product(*args, **kwds)" |
| 906 | pools = [tuple(pool) for pool in args] * repeat |
Raymond Hettinger | fc40b30 | 2020-11-29 10:47:22 -0800 | [diff] [blame] | 907 | return tuple(map(random.choice, pools)) |
Raymond Hettinger | fc91aa2 | 2010-03-28 18:27:13 +0000 | [diff] [blame] | 908 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 909 | def random_permutation(iterable, r=None): |
| 910 | "Random selection from itertools.permutations(iterable, r)" |
| 911 | pool = tuple(iterable) |
| 912 | r = len(pool) if r is None else r |
| 913 | return tuple(random.sample(pool, r)) |
Raymond Hettinger | 063a4b6 | 2010-04-02 04:18:18 +0000 | [diff] [blame] | 914 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 915 | def random_combination(iterable, r): |
| 916 | "Random selection from itertools.combinations(iterable, r)" |
| 917 | pool = tuple(iterable) |
| 918 | n = len(pool) |
| 919 | indices = sorted(random.sample(range(n), r)) |
| 920 | return tuple(pool[i] for i in indices) |
Raymond Hettinger | 063a4b6 | 2010-04-02 04:18:18 +0000 | [diff] [blame] | 921 | |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 922 | def random_combination_with_replacement(iterable, r): |
| 923 | "Random selection from itertools.combinations_with_replacement(iterable, r)" |
| 924 | pool = tuple(iterable) |
| 925 | n = len(pool) |
Raymond Hettinger | fc40b30 | 2020-11-29 10:47:22 -0800 | [diff] [blame] | 926 | indices = sorted(random.choices(range(n), k=r)) |
Raymond Hettinger | 30c7362 | 2010-12-01 23:45:20 +0000 | [diff] [blame] | 927 | return tuple(pool[i] for i in indices) |
Raymond Hettinger | 063a4b6 | 2010-04-02 04:18:18 +0000 | [diff] [blame] | 928 | |
Raymond Hettinger | d37258d | 2018-01-13 10:35:40 -0800 | [diff] [blame] | 929 | def nth_combination(iterable, r, index): |
Raymond Hettinger | fc40b30 | 2020-11-29 10:47:22 -0800 | [diff] [blame] | 930 | "Equivalent to list(combinations(iterable, r))[index]" |
Raymond Hettinger | d37258d | 2018-01-13 10:35:40 -0800 | [diff] [blame] | 931 | pool = tuple(iterable) |
| 932 | n = len(pool) |
| 933 | if r < 0 or r > n: |
| 934 | raise ValueError |
| 935 | c = 1 |
| 936 | k = min(r, n-r) |
| 937 | for i in range(1, k+1): |
| 938 | c = c * (n - k + i) // i |
| 939 | if index < 0: |
| 940 | index += c |
| 941 | if index < 0 or index >= c: |
| 942 | raise IndexError |
| 943 | result = [] |
| 944 | while r: |
| 945 | c, n, r = c*r//n, n-1, r-1 |
| 946 | while index >= c: |
| 947 | index -= c |
| 948 | c, n = c*(n-r)//n, n-1 |
| 949 | result.append(pool[-1-n]) |
| 950 | return tuple(result) |
| 951 | |