Python 模块:itertools

在某些情况下,我们通常需要对序列进行一些复杂的操作,比如从序列中选出一部分元素做排列,组合,笛卡尔积等。如果自己实现这个操作未免太繁琐了,而且还会占用大量的空间,这个时候我们可以求助于 Python 模块——itertools。这个模块总共有 3 部分——无穷迭代器,根据最短输入序列的长度停止的迭代器,排列组合迭代器。

无穷迭代器

常用的无穷迭代器总共有 3 个——count,cycle,repeat,我们先来看一下 count。

count

love

count 顾名思义——数数,这个函数有两个默认参数,第 1 个参数表示从几开始数(默认值 0),第 2 个参数表示每次数的步长(默认值 1)。

import itertools
count = itertools.count()
for c in count:
    print(c)

运行结果输出 0 1 2 3 4 5……它会一直数下去,根本停不下来。

cycle

love

cycle 顾名思义——周期,这个函数有一个必选参数,这个必选参数必须是可迭代的。

import itertools
cycle = itertools.cycle([1, 2, 3])
for c in cycle:
    print(c)

运行结果输出 1 2 3 1 2 3……一直周而复始,永不停息。

repeat

love

repeat 顾名思义——重复,这个函数有一个必选参数,这个必选参数类型是任意的。

import itertools
cycle = itertools.repeat(None)
for c in cycle:
    print(c)

运行结果会一直重复输出 None,也不会停。

根据最短输入序列的长度停止的迭代器

这样的迭代器就比较多了,常用的有 accumulate,chain,chain.from_iterable,compress,dropwhile,filterfalse,islice,starmap,takewhile,ziplongest。

accumulate

love

accumulate 顾名思义——累计,传入一个必选参数和一个可选参数,其中必选参数是一个可迭代对象,可选参数是一个函数(默认值 operator.add)。

import itertools
a = itertools.accumulate([1, 2, 3])
for a0 in a:
    print(a0)

运行结果如图所示。

1 前面没有东西,所以还是 1,2 前面有一个 1,所以 2+1 = 3,3 前面有 2 和 1,所以 3+2+1 = 6。如果我不想做累加,想做累乘怎么办?很简单,指定第二个参数就可以了,就像这样:itertools.accumulate([1, 2, 3], operator.mul)。

chain

love

chain 顾名思义——链条,主要用来把多个序列连在一起做迭代。有些人会认为把多个序列连在一起完全使用操作符 + 就可以实现了,确实可以,但是这特别消耗内存资源,比如下面这种情况。

import itertools
a = [1, 2, 3]
b = [4, 5, 6]
a = a+b

刚开始是两个长度为 3 的列表,当执行完 a = a+b 这一行时,有些人会认为这就仅仅是一个长度为 6 的列表,其实并不是,目前变量 b 所指向的内存并没有被 free,因此内存中还是有 2 个列表,一个长度为 3,一个长度为 6,很明显这比做 + 之前占用的内存要多。为了避免这种情况,我们可以使用 chain。chain 函数参数个数不定,但要求每一个参数都是可迭代对象。

import itertools
chain = itertools.chain([1, 2, 3], [4, 5, 6])
for c in chain:
    print(c)

运行结果如图所示。

chain.from_iterable

love

和 chain 一样,只不过这次不是传入多个参数,而是直接传入一个可迭代对象,如下所示。

import itertools
chain = itertools.chain.from_iterable([[1, 2, 3], [4, 5, 6]])
for c in chain:
    print(c)

运行结果同上。

compress

love

compress 是用来做筛选的,两个必选参数,第一个是待筛选的序列 data,第二个是选择器列表 selectors,如果 selectors[i] 为真,就返回 data[i]。注意:如果 data 和 selectors 长度不相等,以最短的那个为基准!

import itertools
compress = itertools.compress('ABCDEF', [1, 0, 1, 0, 1, 1])
for c in compress:
    print(c)

运行结果如图所示。

dropwhile 和 takewhile

love

这个函数有两个参数,第一个参数是一个函数 pred,第二个参数是一个序列 seq。如果 pred(seq[i]) 为真,直接丢弃,如果 pred(seq[i]) 为假,就不丢弃,注意:当 pred(seq[i]) 为假时,即使 pred(seq[j]) 为真也不会被丢弃(j > i)。

import itertools
d = itertools.dropwhile(lambda x: x < 5, [1, 4, 6, 4, 1])
for d0 in d:
    print(d0)

运行结果如图所示。

很明显,前两个元素被丢弃,后三个元素被保留在迭代器中被迭代输出了。顺便说一下 takewhile,它和 dropwhile 相反,如果这段代码把 dropwhile 改成 takewhile,输出结果就是 1 4,也就是前两个元素。

starmap

love

starmap 是 map 的一个升级版本,比如有这么一个要求,一个列表 a = [(2, 5), (3, 2), (10, 3)],计算 pow(2, 5),pow(3, 2),pow(10, 3)。用 map 实现确实可以,但是 map 每次计算都是不带解包的,所以要想使用 map 必须要多一步解包的操作,显得有些繁琐。而 starmap 给你封装好了,直接用即可。

import itertools
starmap = itertools.starmap(pow, [(2, 5), (3, 2), (10, 3)])
for s in starmap:
    print(s)

运行结果如图所示。

zip_longest

love

zip_longest 是相对于 zip 而言的,这两个函数的第一个参数一样,都是多个可迭代对象。执行的时候都是从每个参数中收集对应的元素,先收集第一个参数,第二个参数……的第一个元素,然后收集第一个参数,第二个参数……的第二个元素,以此类推。那么 zip 和 zip_longest 有什么区别呢?当每个参数的长度不一样时,zip 以最短的为基准,而 zip_longest 以最长的为基准,那么 zip_longest 有缺失怎么办,其实这个函数还有一个默认参数 fillvalue(默认值 None),这个参数就是用来决定有缺失时填充什么。

import itertools
a = itertools.zip_longest('ABCD', 'xy', fillvalue='-')
for a0 in a:
    print(a0)

运行结果如图所示。

排列组合迭代器

排列组合迭代器有 product,permutations,combinations,combinations_with_replacement。

product

love

product 主要用来计算多个序列的笛卡尔积,传入一个可变长参数和一个默认参数 repeat(默认值 1)。这个可变长参数要求每一个都是可迭代对象。

import itertools
x = [1, 2, 3]
y = [4, 5, 6]
a = itertools.product(x, y)
for a0 in a:
    print(a0)

运行结果如图所示。

如果有两个一样的序列 x 做笛卡尔积,我是不是要写成 product(x, x) 呢?完全没有必要,只要使用 repeat 参数就可以了,写成 product(x, repeat=2) 就可以了。

permutations

love

这个函数就是用来做排列的,比如我输入 (1, 2, 3),输出 (1, 2, 3),(1, 3, 2),(2, 1, 3),(2, 3, 1),(3, 1, 2),(3, 2, 1),我就可以使用这个函数,如下所示。

import itertools
a = itertools.permutations([1, 2, 3])
for a0 in a:
    print(a0)

运行结果如图所示。

那么如果我想输出 (1, 2),(1, 3),(2, 1),(2, 3),(3, 1),(3, 2) 该怎么办呢?其实很简单,permutations([1, 2, 3], 2) 就可以了。

combinations 和 combinations_with_replacement

love

这两个函数都是是用来做组合的,从一个序列中选取几个元素,迭代所有的可能。当然这两个是有区别的,其中 combinations 是不能重复选择,而 combinations_with_replacement 是可以重复选择的。

import itertools
print('combinations:')
a = itertools.combinations([1, 2, 3, 4, 5], 2)
for a0 in a:
    print(a0)
print('combinations_with_replacement:')
a = itertools.combinations_with_replacement([1, 2, 3, 4, 5], 2)
for a0 in a:
    print(a0)

运行结果如图所示。

原文发布于微信公众号 - 小陈学Python(gh_a29b1ed16571)

原文发表时间:2019-05-23

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

发表于

我来说两句

0 条评论
登录 后参与评论

扫码关注云+社区

领取腾讯云代金券