itertools模块

itertools: python内置模块，提供了很多关于迭代器相关的方法。使用我们有能力在处理大量数据时，不会大量占用内存，而是像一个管道那样从源头获取一个元素，处理一个元素。

下面是三个构造无限内容的迭代器：

itertools.count(start=0, step=1): 用于构造无限长度的迭代器。而range对象则是有限长度的。除此之外，step可以设为负数。start和step都可以是float类型。一般用于配合其他方法使用。

 # count(10) → 10 11 12 13 14 ...
 # count(2.5, 0.5) → 2.5 3.0 3.5 ...

itertools.cycle(iterable): 对一个迭代器(通常是有限迭代器)的无限循环迭代。

# cycle('ABCD') → A B C D A B C D A B C D ...

itertools.repeat(object[, times]): 返回一个迭代器，重复一个object，可以指定循环次数，否则就是无限循环。

from itertools import repeat

list(repeat([1, 2, 3], 3))  # [[1, 2, 3], [1, 2, 3], [1, 2, 3]]
list(repeat(10, 5))  # [10, 10, 10, 10, 10]

下面是常用迭代器函数:

itertools.accumulate(iterable[, function, *, initial=None]): 接收一个可迭代对象和一个函数，将函数作用在每个元素上的结果以迭代器的形式返回。 如果不传递function参数，则默认为operator.add。

 # accumulate([1,2,3,4,5]) → 1 3 6 10 15
 # accumulate([1,2,3,4,5], initial=100) → 100 101 103 106 110 115
 # accumulate([1,2,3,4,5], operator.mul) → 1 2 6 24 120

可以看出，其余functools.reduce函数的不同在于，reduce函数返回的是最终结果，相当于accumulate返回的迭代器的最后一个元素。

itertools.batched(iterable, n, *, strict=False): 将一个可迭代对象按照n个元素一组，并以迭代器的形式返回，每个迭代器元素是一个n元组。最后一个元素的长度可能小于n。 但是当strict=True时，如果最后一个长度小于n会抛异常。

# batched('ABCDEFG', 3) → ABC DEF G

由于batched函数是在3.12版本里新加的，下面写一个模拟的实现。但是效果不是完全等价的。最后一个元组使用了None来填充。

from itertools import zip_longest

lst = 'ABCDEFG'

def custom_batched(iterable, n):
    it = iter(iterable)
    return zip_longest(*[it] * n)

print(list(custom_batched(lst, 3)))  # [('A', 'B', 'C'), ('D', 'E', 'F'), ('G', None, None)]

itertools.chain(*iterables): 返回一个迭代器，相当于把多个可迭代对象串联起来，也相当于下面的等价实现。

def chain(*iterables):
    # chain('ABC', 'DEF') → A B C D E F
    for iterable in iterables:
        yield from iterable

chain.from_iterable(iterable): 返回一个迭代器，相当于把嵌套的可迭代对象展平，并将它们连接成一个大的可迭代对象。

def from_iterable(iterables):
    # chain.from_iterable(['ABC', 'DEF']) → A B C D E F
    for iterable in iterables:
        yield from iterable

itertools.compress(data, selectors): 返回一个迭代器，每个元素来自data中，并且在相同位置的selectors中对应的值为真。
类似于下面的实现：

def compress(data, selectors):
    # compress('ABCDEF', [1,0,1,0,1,1]) → A C E F
    return (datum for datum, selector in zip(data, selectors) if selector)

使用示例：

from itertools import compress

iters = 'ABCDEFG'
selectors = [True, True, False, False, True]

print(list(compress(iters, selectors))) # ['A', 'B', 'E']

itertools.starmap(function, iterable): 返回一个迭代器，每个元素为将iterable里每个元素做为参数传给function后的执行结果。
与内置函数map()非常接近，只不过，如果function接收多个参数，iterable里的每个元素应该为一个多元素构成的元组。
等价实现类似于下面函数：

def starmap(function, iterable):
    # starmap(pow, [(2,5), (3,2), (10,3)]) → 32 9 1000
    for args in iterable:
        yield function(*args)

itertools.takewhile(predicate, iterable): 返回一个迭代器，元素来自iterable，当predicate函数作用于iterable为False时，迭代器元素结束。

from itertools import takewhile

lst = [0,1,2,4,5,1,2,0]
print(list(takewhile(lambda x: x < 4, lst)))  # [0, 1, 2]

从上面示例可以看出，即使列表中元素4,5后面仍要满足条件的元素，但返回的迭代器已经停止了。

itertools.dropwhile(predicate, iterable): 跟takewhile函数正好相反，当predicate作用于iterable元素返回False时，才开始将之与后面的元素以迭代器的形式返回。

from itertools import dropwhile

lst = [0,1,2,4,5,1,2,0]

print(list(dropwhile(lambda x: x < 3, lst)))  # [4, 5, 1, 2, 0]

itertools.filterfalse(predicate, iterable): 将predicate作用于iterable中的每个元素，所果返回为False，则这个元素将被包括在返回的迭代器里。 如果predicate是None，则相当于是传了一个bool()函数。

from itertools import filterfalse

lst = [0,1,2,4,5,1,2,0]

print(list(filterfalse(lambda x: x < 3, lst)))  # [4, 5]
print(list(filterfalse(lambda x: x > 3, lst)))  # [0, 1, 2, 1, 2, 0]
print(list(filterfalse(None, lst)))  # [0, 0]

itertools.pairwise(iterable): 将迭代器中每个元素都跟它的后一个元素做为一个元组来返回。总的数量将比原可迭代对象少1。

from itertools import pairwise

print(list(pairwise('ABCDEFG')))  # [('A', 'B'), ('B', 'C'), ('C', 'D'), ('D', 'E'), ('E', 'F'), ('F', 'G')]

此函数大致相当于下面的模拟实现：

def pairwise(iterable):
    # pairwise('ABCDEFG') → AB BC CD DE EF FG

    iterator = iter(iterable)
    a = next(iterator, None)

    for b in iterator:
        yield a, b
        a = b

itertools.zip_longest(*iterables, fillvalue=None): 行为与内置函数zip(*iterables, strict=False)大致相同，只不过当各个iterable长度不同时，由最长的iterable来决定长度，元组其他位置的值用fillvalue参数的值提供补充。

from itertools import zip_longest
    
for item in zip_longest([1, 2, 3], "abcd", fillvalue=9):
    print(item)

输出结果：

(1, 'a')
(2, 'b')
(3, 'c')
(9, 'd')

itertools.tee(iterable, n=2): 从一个可迭代对象里返回n个独立的迭代器。n默认为2。

from itertools import tee

it1, it2 = tee([1, 2, 3, 4, 5, 6], 2)

for item in it1:
    print(item)

print('-' * 25, '分隔线', '-' * 25)

for item in it2:
    print(item)

输出结果：

1
2
3
4
5
6
------------------------- 分隔线 -------------------------
1
2
3
4
5
6

可以看到，第一个迭代器耗尽后，不会影响第二个迭代器。
然而tee(iterable, n=2)并不等价于 it1 = iter(iterable), it2 = iter(iterable), 而是比其有更高的内存性能，其内部实现并非这么简单。 在大数据集上生成独立迭代器，推荐使用tee函数，而不是多次使用iter(iterable)的方法。

itertools.islice: 返回对一个可迭代对象进行切片后的元素的迭代器。有如下两种函数签名。
itertools.islice(iterable, stop): 返回stop下标位置之前的元素的迭代器。

from itertools import islice

lst = [1,2,3,4,5,6,7]

print(list(islice(lst, 4)))  # [1, 2, 3, 4]

itertools.islice(iterable, start, stop[, step]): 返回(start, stop)之间的元素的迭代器，是左闭右开区间。

from itertools import islice

lst = [1,2,3,4,5,6,7]

print(list(islice(lst, 3, 6)))  # [4, 5, 6]

注意，使用islice对同一个迭代器进行切片，创建多个切片迭代器后，它们之间会互相影响。看下面的例子:

from itertools import islice
import string

deck = string.ascii_lowercase
itr = iter(deck)
n = 4

print('-' * 25, '分隔线', '-' * 25)
for item in zip(islice(itr, n), islice(itr, n), islice(itr, n)):
    print(item)

输出结果：

------------------------- 分隔线 -------------------------
('a', 'b', 'c')
('d', 'e', 'f')
('g', 'h', 'i')
('j', 'k', 'l')

上面示例中，每个切片迭代器的长度都是4，所以zip最终返回了4个元组。由于都基于同一个迭代器做的切片，每个切片迭代器迭代一个新元素后，原始迭代器也会同步向前递进一个元素。并且是相同的元素。后面的切片迭代器再迭代时，就会递进到原始迭代器的下一个元素，因此以第一个切片迭代器为例，它其实返回的是a,d,g,j这4个元素。

那么我们换一种写法，看看结果是否相同：

from itertools import islice
import string

deck = string.ascii_lowercase
itr = iter(deck)
n = 4

print('-' * 25, '分隔线', '-' * 25)
for item in zip(*[islice(itr, n)] * 3):
    print(item)

输出结果：

------------------------- 分隔线 -------------------------
('a', 'b', 'c')

从示例可以看到，这次只输出了一个元组，并且只是原始迭代器中的前三个元素，这是为什么呢？
原因是[islice(itr, n)] * 3 只会创建一个切片迭代器，在列表中相当于三个元素同时指向一个对象，而这个切片迭代器长度为4，当在第一轮迭代时就被递进了三个元素，第二轮迭代时会导致不够再迭代出三个元素，由于zip函数的特性，以最短的可迭代对象来决定返回元组的个数，所以只够返回第一个元组，切片迭代器就被耗尽了。
那么前一种方式多次调用islice(itr, n)产生的是三个不同的切片迭代器么？ 是的，但我们先打印出来看看:

from itertools import islice
import string

deck = string.ascii_lowercase
itr = iter(deck)
n = 4

print('-' * 25, '分隔线', '-' * 25)
print(islice(itr, n))
print(islice(itr, n))
print(islice(itr, n))

输出结果：

------------------------- 分隔线 -------------------------
<itertools.islice object at 0x0000027A42C8A930>
<itertools.islice object at 0x0000027A42C8A930>
<itertools.islice object at 0x0000027A42C8A930>

奇怪，输出显示三个切片迭代器指向了同一个内存地址，这不应该是同一个对象么?
对此，AI做了如下解释:
Python 可能会优化对象的创建，当某个对象短暂存在，而在 print() 语句执行完后马上就被销毁，Python 可能会在下一个对象分配时重用相同的内存地址。
也就是说，islice(itr, n) 迭代器对象被 print() 打印后，它的引用立即释放，所以新的 islice 可能会获得相同的内存地址，导致 print() 显示的内容一致。
但如果你存储这些 islice 迭代器，就能发现它们确实是不同的对象。

from itertools import islice
import string

deck = string.ascii_lowercase
itr = iter(deck)
n = 4

print('-' * 25, '分隔线', '-' * 25)
print(islice(itr, n))
print(islice(itr, n))
print(islice(itr, n))

print('-' * 25, '分隔线', '-' * 25)
a = islice(itr, n)
b = islice(itr, n)
c = islice(itr, n)

print(a)
print(b)
print(c)

print('-' * 25, '分隔线', '-' * 25)
print((islice(itr, n) is islice(itr, n)))

输出结果：

------------------------- 分隔线 -------------------------
<itertools.islice object at 0x000001C9712BA9D0>
<itertools.islice object at 0x000001C9712BA9D0>
<itertools.islice object at 0x000001C9712BA9D0>
------------------------- 分隔线 -------------------------
<itertools.islice object at 0x000001C9712BA9D0>
<itertools.islice object at 0x000001C9712BAA20>
<itertools.islice object at 0x000001C97130C8B0>
------------------------- 分隔线 -------------------------
False

这种现象有时还体现为正常跑和debug模式看到的效果不一样。 debug模式运行时看到不是相同的对象地址，再正常运行时打印出来的却显示是相同的对象地址。

itertools.groupby(iterable, key=None): 用于将可迭代对象中的连续相同元素分组。它返回一个迭代器，生成由键和对应分组组成的元组 (key, group)，这里的group自身又是一个迭代器。

from itertools import groupby

for k, g in groupby('AAAABBBCCDAABBB'):
    print('-' * 25, '分隔线', '-' * 25)
    print(f'{k=}')
    for item in g:
        print(item)

print('-' * 25, '分隔线', '-' * 25)
print([list(g) for k, g in groupby('AAAABBBCCDAABBB')])

输出结果：

------------------------- 分隔线 -------------------------
k='A'
A
A
A
A
------------------------- 分隔线 -------------------------
k='B'
B
B
B
------------------------- 分隔线 -------------------------
k='C'
C
C
------------------------- 分隔线 -------------------------
k='D'
D
------------------------- 分隔线 -------------------------
k='A'
A
A
------------------------- 分隔线 -------------------------
k='B'
B
B
B
------------------------- 分隔线 -------------------------
[['A', 'A', 'A', 'A'], ['B', 'B', 'B'], ['C', 'C'], ['D'], ['A', 'A'], ['B', 'B', 'B']]

由于其分组是依赖于连续的元素，这点与sql里的group by不同。 所以如果想对全局元素分组，要确保先排序。

from itertools import groupby

words = ['apple', 'banana', 'apricot', 'blueberry', 'avocado']
sorted_words = sorted(words, key=lambda w: w[0])  # 按首字母排序

for key, group in groupby(sorted_words, key=lambda w: w[0]):
    print(key, list(group))

输出结果：

a ['apple', 'apricot', 'avocado']
b ['banana', 'blueberry']

下面是排列组合相关函数

itertools.product(*iterables, repeat=1): 求多个可迭代对象之间组合的迪卡尔积。即所有可能的组合，但各个组合内元素的顺序保留了多个iterables之间的相对顺序。

from itertools import product

for item in product("ab",[1,2]):
    print(item)

print('-' * 25, '分隔线', '-' * 25)
for item in product(range(2), repeat=3):
    print(item)

print('-' * 25, '分隔线', '-' * 25)
for item in product("ab",[1,2], repeat=2):
    print(item)

print('-' * 25, '分隔线', '-' * 25)
for item in product("ab",[1,2], "ab", [1,2]):
    print(item)

注意，这个函数会将每个位置看做是不同的值，所以如果可迭代对象中本身有不同的值，笛卡尔积的结果中就会有相同的值。

from itertools import product

for item in product("aa",[1,2]):
    print(item)

输出结果：

('a', 1)
('a', 2)
('a', 1)
('a', 2)

itertools.permutations(iterable, r=None) 从一个可迭代对象中取出一定数量的元素，有多少种取法，不同的顺序算不同的取法。 这个相当于数学中的排列概念，从n个元素中取出m个进行排列，数学公式为 n!/(n-m)!，即n的阶乘除以(n-m)的阶乘。

from itertools import permutations

for item in permutations([1,2,3,4], 2):
    print(item)

print('-' * 25, '分隔线', '-' * 25)

for item in permutations([1,2,3,4]):
    print(item)

输出结果：

(1, 2)
(1, 3)
(1, 4)
(2, 1)
(2, 3)
(2, 4)
(3, 1)
(3, 2)
(3, 4)
(4, 1)
(4, 2)
(4, 3)
------------------------- 分隔线 -------------------------
(1, 2, 3, 4)
(1, 2, 4, 3)
(1, 3, 2, 4)
(1, 3, 4, 2)
(1, 4, 2, 3)
(1, 4, 3, 2)
(2, 1, 3, 4)
(2, 1, 4, 3)
(2, 3, 1, 4)
(2, 3, 4, 1)
(2, 4, 1, 3)
(2, 4, 3, 1)
(3, 1, 2, 4)
(3, 1, 4, 2)
(3, 2, 1, 4)
(3, 2, 4, 1)
(3, 4, 1, 2)
(3, 4, 2, 1)
(4, 1, 2, 3)
(4, 1, 3, 2)
(4, 2, 1, 3)
(4, 2, 3, 1)
(4, 3, 1, 2)
(4, 3, 2, 1)

注意，这个函数会将每个位置看做是不同的值，去排列。所以如果可迭代对象中本身有不同的值，排列结果中就会有相同的值。

from itertools import permutations

for item in permutations([1,2,2,3], 2):
    print(item)

输出结果：

(1, 2)
(1, 2)
(1, 3)
(2, 1)
(2, 2)
(2, 3)
(2, 1)
(2, 2)
(2, 3)
(3, 1)
(3, 2)
(3, 2)

itertools.combinations(iterable, r)： 返回一个迭代器，对应从iterable中取出r个元素的各种组合，不考虑顺序。 对应数学中的组合的概念。从n个元素中取出m个元素的组合数量为 m!/n!/(m-n)!。

from itertools import product, combinations

for item in combinations('ABCD', 2):
    print(item)

输出结果：

('A', 'B')
('A', 'C')
('A', 'D')
('B', 'C')
('B', 'D')
('C', 'D')

组合数量同样是按照元素位置计算的，不同位置看做不同的元素，所以当原始iterable本身包含重复值时，返回的组合也有重复值。

from itertools import product, combinations

for item in combinations('ABBD', 2):
    print(item)

输出结果：

('A', 'B')
('A', 'B')
('A', 'D')
('B', 'B')
('B', 'D')
('B', 'D')

itertools.combinations_with_replacement(iterable, r): 与combinations功能类似，也是生成组合，而不是排列。但iterable中的每个元素都可以重复使用r次。例如:combinations_with_replacement('ABC', 2) → AA AB AC BB BC CC，这里有AB而没有BA。元组里元素顺序类似数学中的打枪法，所以先有AB，就不会再有BA这个组合了。

from itertools import combinations, combinations_with_replacement

for item in combinations_with_replacement([1, 2, 3], 3):
    print(item)

输出结果：

(1, 1, 1)
(1, 1, 2)
(1, 1, 3)
(1, 2, 2)
(1, 2, 3)
(1, 3, 3)
(2, 2, 2)
(2, 2, 3)
(2, 3, 3)
(3, 3, 3)

注意，函数仍然会把iterable中每个位置的元素当成是不同元素进行组合，请看下面示例：

from itertools import combinations, combinations_with_replacement

for item in combinations_with_replacement('ABB', 2):
    print(item)

输出结果：

('A', 'A')
('A', 'B')
('A', 'B')
('B', 'B')
('B', 'B')
('B', 'B')