python itertools module

发表于 分类于 ，

简介

itertools 是python的迭代器模块，itertools提供的工具相当高效且节省内存

用法

1
2
3
>>> import itertools
>>> dir(itertools)
['__doc__', '__loader__', '__name__', '__package__', '__spec__', '_grouper', '_tee', '_tee_dataobject', 'accumulate', 'chain', 'combinations', 'combinations_with_replacement', 'compress', 'count', 'cycle', 'dropwhile', 'filterfalse', 'groupby', 'islice', 'permutations', 'product', 'repeat', 'starmap', 'takewhile', 'tee', 'zip_longest']
  • 无穷迭代器
迭代器 实参 结果 示例
count() start, [step] start, start+step, start+2*step, … count(10) --> 10 11 12 13 14...
cycle() p p0, p1, … plast, p0, p1, … cycle('ABCD') --> A B C D A BC D ...
repeat() elem [,n] elem, elem, elem, … 重复无限次或n次 repeat(10, 3) --> 10 10 10
  • 根据最短输入序列长度停止的迭代器
迭代器 实参 结果 示例
accumulate() p [,func] p0, p0+p1, p0+p1+p2, … accumulate([1,2,3,4,5]) --> 1 36 10 15
chain() p, q, … p0, p1, … plast, q0, q1, … chain('ABC', 'DEF') --> A B C DE F
chain.from_iterable() iterable – 可迭代对象 p0, p1, … plast, q0, q1, … chain.from_iterable(['ABC','DEF']) --> A B C D E F
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
dropwhile() pred, seq seq[n], seq[n+1], … 从pred首次真值测试失败开始 dropwhile(lambda x: x<5,[1,4,6,4,1]) --> 6 4 1
filterfalse() pred, seq seq中pred(x)为假值的元素，x是seq中的元素。 filterfalse(lambda x: x%2,range(10)) --> 0 2 4 6 8
groupby() iterable[, key] 根据key(v)值分组的迭代器
islice() seq, [start,] stop [, step] seq[start:stop:step]中的元素 islice('ABCDEFG', 2, None) --> CD E F G
starmap() func, seq func(seq[0]), func(seq[1]), … starmap(pow, [(2,5), (3,2),(10,3)]) --> 32 9 1000
takewhile() pred, seq seq[0], seq[1], …, 直到pred真值测试失败 takewhile(lambda x: x<5,[1,4,6,4,1]) --> 1 4
tee() it, n it1, it2, … itn 将一个迭代器拆分为n个迭代器
zip_longest() p, q, … (p[0], q[0]), (p[1], q[1]), … zip_longest('ABCD', 'xy',fillvalue='-') --> Ax By C- D-
  • 排列组合迭代器
迭代器 实参 结果
product() p, q, … [repeat=1] 笛卡尔积，相当于嵌套的for循环
permutations() p[, r] 长度r元组，所有可能的排列，无重复元素
combinations() p, r 长度r元组，有序，无重复元素
combinations_with_replacement() p, r 长度r元组，有序，元素可重复
product('ABCD', repeat=2) AA AB AC AD BA BB BC BD CA CB CC CD DA DBDC DD
permutations('ABCD', 2) AB AC AD BA BC BD CA CB CD DA DB DC
combinations('ABCD', 2) AB AC AD BC BD CD
combinations_with_replacement('ABCD',2) AA AB AC AD BB BC BD CC CD DD

itertools.count(start=0, step=1)

创建一个迭代器，它从 start 值开始，返回均匀间隔的值。常用于 map() 中的实参来生成连续的数据点。此外，还用于 zip() 来添加序列号。大致相当于：

1
2
3
4
5
6
7
def count(start=0, step=1):
    # count(10) --> 10 11 12 13 14 ...
    # count(2.5, 0.5) -> 2.5 3.0 3.5 ...
    n = start
    while True:
        yield n
        n += step

当对浮点数计数时，替换为乘法代码有时精度会更好，例如： (start + step * i for i in count()) 。

在 3.1 版更改: 增加参数 step ，允许非整型。

1
2
3
4
>>> import itertools

>>> list(zip("abcd", itertools.count()))
[('a', 0), ('b', 1), ('c', 2), ('d', 3)]

itertools.cycle(iterable)

创建一个迭代器，返回 iterable 中所有元素并保存一个副本。当取完 iterable 中所有元素，返回副本中的所有元素。无限重复。大致相当于：

1
2
3
4
5
6
7
8
9
def cycle(iterable):
    # cycle('ABCD') --> A B C D A B C D A B C D ...
    saved = []
    for element in iterable:
        yield element
        saved.append(element)
    while saved:
        for element in saved:
              yield element

注意，该函数可能需要相当大的辅助空间（取决于 iterable 的长度）。

itertools.repeat(object[, times])

创建一个迭代器，不断重复 object 。除非设定参数 times ，否则将无限重复。可用于 map() 函数中的参数，被调用函数可得到一个不变参数。也可用于 zip() 的参数以在元组记录中创建一个不变的部分。

大致相当于：

1
2
3
4
5
6
7
8
def repeat(object, times=None):
    # repeat(10, 3) --> 10 10 10
    if times is None:
        while True:
            yield object
    else:
        for i in range(times):
            yield object

repeat 最常见的用途就是在 map 或 zip 提供一个常量流：

1
2
>>> list(map(pow, range(10), repeat(2)))
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

Itertool函数

itertools.accumulate(iterable[, func])

创建一个迭代器，返回累加和或其他二元函数的累加结果（通过可选参数 func 指定）。如果提供了 func ，它应是2个参数的函数。输入 iterable 元素类型应是 func 能支持的任意类型。（例如，对于默认的加法操作，元素可以是任一支持加法的类型，包括 Decimal 或 Fraction ）。如果可迭代对象的输入为空，输出也为空。

大致相当于：

1
2
3
4
5
6
7
8
9
10
11
12
13
def accumulate(iterable, func=operator.add):
    'Return running totals'
    # accumulate([1,2,3,4,5]) --> 1 3 6 10 15
    # accumulate([1,2,3,4,5], operator.mul) --> 1 2 6 24 120
    it = iter(iterable)
    try:
        total = next(it)
    except StopIteration:
        return
    yield total
    for element in it:
        total = func(total, element)
        yield total

func 参数有几种用法。它可以被设为 min() 最终得到一个最小值，或者设为 max() 最终得到一个最大值，或设为 operator.mul() 最终得到一个乘积。摊销表可通过累加利息和支付款项得到。给iterable设置初始值并只将参数 func 设为累加总数可以对一阶 递归关系 建模。

函数 functools.reduce() ，它只返回一个最终累积值

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
>>> data = [3, 4, 6, 2, 1, 9, 0, 7, 5, 8]
>>> list(accumulate(data, operator.mul))     # running product
[3, 12, 72, 144, 144, 1296, 0, 0, 0, 0]
>>> list(accumulate(data, max))              # running maximum
[3, 4, 6, 6, 6, 9, 9, 9, 9, 9]

# Amortize a 5% loan of 1000 with 4 annual payments of 90
>>> cashflows = [1000, -90, -90, -90, -90]
>>> list(accumulate(cashflows, lambda bal, pmt: bal*1.05 + pmt))
[1000, 960.0, 918.0, 873.9000000000001, 827.5950000000001]

# Chaotic recurrence relation https://en.wikipedia.org/wiki/Logistic_map
>>> logistic_map = lambda x, _:  r * x * (1 - x)
>>> r = 3.8
>>> x0 = 0.4
>>> inputs = repeat(x0, 36)     # only the initial value is used
>>> [format(x, '.2f') for x in accumulate(inputs, logistic_map)]
['0.40', '0.91', '0.30', '0.81', '0.60', '0.92', '0.29', '0.79', '0.63',
 '0.88', '0.39', '0.90', '0.33', '0.84', '0.52', '0.95', '0.18', '0.57',
 '0.93', '0.25', '0.71', '0.79', '0.63', '0.88', '0.39', '0.91', '0.32',
 '0.83', '0.54', '0.95', '0.20', '0.60', '0.91', '0.30', '0.80', '0.60']

itertools.chain(*iterables)

创建一个迭代器，它首先返回第一个可迭代对象中所有元素，接着返回下一个可迭代对象中所有元素，直到耗尽所有可迭代对象中的元素。可将多个序列处理为单个序列。大致相当于：

1
2
3
4
5
def chain(*iterables):
    # chain('ABC', 'DEF') --> A B C D E F
    for it in iterables:
        for element in it:
            yield element

classmethod chain.from_iterable(iterable)

构建类似 chain() 迭代器的另一个选择。从一个单独的可迭代参数中得到链式输入，该参数是延迟计算的。大致相当于：

1
2
3
4
5
def from_iterable(iterables):
    # chain.from_iterable(['ABC', 'DEF']) --> A B C D E F
    for it in iterables:
        for element in it:
            yield element
1
2
>>> list(itertools.chain().from_iterable(["acd", "123"]))  
['a', 'c', 'd', '1', '2', '3']

itertools.compress(data, selectors)

创建一个迭代器，它返回 data 中经 selectors 真值测试为 True 的元素。迭代器在两者较短的长度处停止。大致相当于：

1
2
3
def compress(data, selectors):
    # compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F
    return (d for d, s in zip(data, selectors) if s)

itertools.dropwhile(predicate, iterable)

创建一个迭代器，如果 predicate 为true，迭代器丢弃这些元素，然后返回其他元素。注意，迭代器在 predicate 首次为false之前不会产生任何输出，所以可能需要一定长度的启动时间。大致相当于：

1
2
3
4
5
6
7
8
9
def dropwhile(predicate, iterable):
    # dropwhile(lambda x: x<5, [1,4,6,4,1]) --> 6 4 1
    iterable = iter(iterable)
    for x in iterable:
        if not predicate(x):
            yield x
            break
    for x in iterable:
        yield x

itertools.filterfalse(predicate, iterable)

创建一个迭代器，只返回 iterable 中 predicate 为 False 的元素。如果 predicate 是 None，返回真值测试为false的元素。大致相当于：

1
2
3
4
5
6
7
def filterfalse(predicate, iterable):
    # filterfalse(lambda x: x%2, range(10)) --> 0 2 4 6 8
    if predicate is None:
        predicate = bool
    for x in iterable:
        if not predicate(x):
            yield x

itertools.groupby(iterable, key=None)

创建一个迭代器，返回 iterable 中连续的键和组。key 是一个计算元素键值函数。如果未指定或为 None，key 缺省为恒等函数（identity function），返回元素不变。一般来说，iterable 需用同一个键值函数预先排序。

groupby() 操作类似于Unix中的 uniq。当每次 key 函数产生的键值改变时，迭代器会分组或生成一个新组（这就是为什么通常需要使用同一个键值函数先对数据进行排序）。这种行为与SQL的GROUP BY操作不同，SQL的操作会忽略输入的顺序将相同键值的元素分在同组中。

返回的组本身也是一个迭代器，它与 groupby() 共享底层的可迭代对象。因为源是共享的，当 groupby() 对象向后迭代时，前一个组将消失。因此如果稍后还需要返回结果，可保存为列表：

1
2
3
4
5
6
groups = []
uniquekeys = []
data = sorted(data, key=keyfunc)
for k, g in groupby(data, keyfunc):
    groups.append(list(g))      # Store group iterator as a list
    uniquekeys.append(k)
1
2
3
4
5
6
7
8
import random

lst = [random.randint(20) for x in range(10)]

def fn(num):
    return ">10" if num > 10 else "<=10"

rst = [(k, list(v)) for k, v initertools.groupby(sorted(lst), key=fn)]

itertools.islice(iterable, stop)

itertools.islice(iterable, start, stop[, step])

创建一个迭代器，返回从 iterable 里选中的元素。如果 start 不是0，跳过 iterable 中的元素，直到到达 start 这个位置。之后迭代器连续返回元素，除非 step 设置的值很高导致被跳过。如果 stop 为 None，迭代器耗光为止；否则，在指定的位置停止。与普通的切片不同，islice() 不支持将 start ， stop ，或 step 设为负值。可用来从内部数据结构被压平的数据中提取相关字段（例如一个多行报告，它的名称字段出现在每三行上）。大致相当于：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
def islice(iterable, *args):
    # islice('ABCDEFG', 2) --> A B
    # islice('ABCDEFG', 2, 4) --> C D
    # islice('ABCDEFG', 2, None) --> C D E F G
    # islice('ABCDEFG', 0, None, 2) --> A C E G
    s = slice(*args)
    start, stop, step = s.start or 0, s.stop or sys.maxsize, s.step or 1
    it = iter(range(start, stop, step))
    try:
        nexti = next(it)
    except StopIteration:
        # Consume *iterable* up to the *start* position.
        for i, element in zip(range(start), iterable):
            pass
        return
    try:
        for i, element in enumerate(iterable):
            if i == nexti:
                yield element
                nexti = next(it)
    except StopIteration:
        # Consume to *stop*.
        for i, element in zip(range(i + 1, stop), iterable):
            pass

如果 start 为 None，迭代从0开始。如果 step 为 None ，步长缺省为1。

itertools.starmap(function, iterable)

创建一个迭代器，使用从可迭代对象中获取的参数来计算该函数。当参数对应的形参已从一个单独可迭代对象组合为元组时（数据已被“预组对”）可用此函数代替 map()。map() 与 starmap() 之间的区别可以类比 function(a,b) 与 function(*c) 的区别。大致相当于:

1
2
3
4
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)

创建一个迭代器，只要 predicate 为真就从可迭代对象中返回元素。大致相当于:

1
2
3
4
5
6
7
def takewhile(predicate, iterable):
    # takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4
    for x in iterable:
        if predicate(x):
            yield x
        else:
            break

itertools.tee(iterable, n=2)

从一个可迭代对象中返回 n 个独立的迭代器。

下面的Python代码能帮助解释 tee 做了什么（尽管实际的实现更复杂，而且仅使用了一个底层的 FIFO 队列）。

大致相当于：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
def tee(iterable, n=2):
    it = iter(iterable)
    deques = [collections.deque() for i in range(n)]
    def gen(mydeque):
        while True:
            if not mydeque:             # when the local deque is empty
                try:
                    newval = next(it)   # fetch a new value and
                except StopIteration:
                    return
                for d in deques:        # load it to all the deques
                    d.append(newval)
            yield mydeque.popleft()
    return tuple(gen(d) for d in deques)

一旦 tee() 实施了一次分裂，原有的 iterable 不应再被使用；否则tee对象无法得知 iterable 可能已向后迭代。

tee 迭代器不是线程安全的。当同时使用由同一个 tee() 调用所返回的迭代器时可能引发 RuntimeError，即使原本的 iterable 是线程安全的。

该迭代工具可能需要相当大的辅助存储空间（这取决于要保存多少临时数据）。通常，如果一个迭代器在另一个迭代器开始之前就要使用大部份或全部数据，使用 list() 会比 tee() 更快。

itertools.zip_longest(iterables, fillvalue=None)

创建一个迭代器，从每个可迭代对象中收集元素。如果可迭代对象的长度未对齐，将根据 fillvalue 填充缺失值。迭代持续到耗光最长的可迭代对象。大致相当于：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
def zip_longest(*args, fillvalue=None):
    # zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-
    iterators = [iter(it) for it in args]
    num_active = len(iterators)
    if not num_active:
        return
    while True:
        values = []
        for i, it in enumerate(iterators):
            try:
                value = next(it)
            except StopIteration:
                num_active -= 1
                if not num_active:
                    return
                iterators[i] = repeat(fillvalue)
                value = fillvalue
            values.append(value)
        yield tuple(values)

如果其中一个可迭代对象有无限长度，zip_longest() 函数应封装在限制调用次数的场景中（例如 islice() 或 takewhile()）。除非指定， fillvalue 默认为 None 。

排列组合迭代器：

itertools.product(*iterables, repeat=1)

可迭代对象输入的笛卡儿积。

大致相当于生成器表达式中的嵌套循环。例如， product(A, B) 和 ((x,y) for x in A for y in B) 返回结果一样。

嵌套循环像里程表那样循环变动，每次迭代时将最右侧的元素向后迭代。这种模式形成了一种字典序，因此如果输入的可迭代对象是已排序的，笛卡尔积元组依次序发出。

要计算可迭代对象自身的笛卡尔积，将可选参数 repeat 设定为要重复的次数。例如，product(A, repeat=4) 和 product(A, A, A, A) 是一样的。

该函数大致相当于下面的代码，只不过实际实现方案不会在内存中创建中间结果。

1
2
3
4
5
6
7
8
9
def product(*args, repeat=1):
    # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy
    # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111
    pools = [tuple(pool) for pool in args] * repeat
    result = [[]]
    for pool in pools:
        result = [x+[y] for x in result for y in pool]
    for prod in result:
        yield tuple(prod)

itertools.permutations(iterable, r=None)

连续返回由 iterable 元素生成长度为 r 的排列。

如果 r 未指定或为 None ，r 默认设置为 iterable 的长度，这种情况下，生成所有全长排列。

排列依字典序发出。因此，如果 iterable 是已排序的，排列元组将有序地产出。

即使元素的值相同，不同位置的元素也被认为是不同的。如果元素值都不同，每个排列中的元素值不会重复。

大致相当于：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
def permutations(iterable, r=None):
    # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
    # permutations(range(3)) --> 012 021 102 120 201 210
    pool = tuple(iterable)
    n = len(pool)
    r = n if r is None else r
    if r > n:
        return
    indices = list(range(n))
    cycles = list(range(n, n-r, -1))
    yield tuple(pool[i] for i in indices[:r])
    while n:
        for i in reversed(range(r)):
            cycles[i] -= 1
            if cycles[i] == 0:
                indices[i:] = indices[i+1:] + indices[i:i+1]
                cycles[i] = n - i
            else:
                j = cycles[i]
                indices[i], indices[-j] = indices[-j], indices[i]
                yield tuple(pool[i] for i in indices[:r])
                break
        else:
            return

permutations() 的代码也可被改写为 product() 的子序列，只要将含有重复元素（来自输入中同一位置的）的项排除。

1
2
3
4
5
6
7
def permutations(iterable, r=None):
    pool = tuple(iterable)
    n = len(pool)
    r = n if r is None else r
    for indices in product(range(n), repeat=r):
        if len(set(indices)) == r:
            yield tuple(pool[i] for i in indices)

当 0 <= r <= n ，返回项个数为 n! / (n-r)! ；当 r > n ，返回项个数为0。

itertools.combinations(iterable, r)

返回由输入 iterable 中元素组成长度为 r 的子序列。

组合按照字典序返回。所以如果输入 iterable 是有序的，生成的组合元组也是有序的。

即使元素的值相同，不同位置的元素也被认为是不同的。如果元素各自不同，那么每个组合中没有重复元素。

大致相当于：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
def combinations(iterable, r):
    # combinations('ABCD', 2) --> AB AC AD BC BD CD
    # combinations(range(4), 3) --> 012 013 023 123
    pool = tuple(iterable)
    n = len(pool)
    if r > n:
        return
    indices = list(range(r))
    yield tuple(pool[i] for i in indices)
    while True:
        for i in reversed(range(r)):
            if indices[i] != i + n - r:
                break
        else:
            return
        indices[i] += 1
        for j in range(i+1, r):
            indices[j] = indices[j-1] + 1
        yield tuple(pool[i] for i in indices)

combinations() 的代码可被改写为 permutations() 过滤后的子序列，（相对于元素在输入中的位置）元素不是有序的。

1
2
3
4
5
6
def combinations(iterable, r):
    pool = tuple(iterable)
    n = len(pool)
    for indices in permutations(range(n), r):
        if sorted(indices) == list(indices):
            yield tuple(pool[i] for i in indices)

当 0 <= r <= n 时，返回项的个数是 n! / r! / (n-r)!；当 r > n 时，返回项个数为0。

itertools.combinations_with_replacement(iterable, r)

返回由输入 iterable 中元素组成的长度为 r 的子序列，允许每个元素可重复出现。

组合按照字典序返回。所以如果输入 iterable 是有序的，生成的组合元组也是有序的。

不同位置的元素是不同的，即使它们的值相同。因此如果输入中的元素都是不同的话，返回的组合中元素也都会不同。

大致相当于：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
def combinations_with_replacement(iterable, r):
    # combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC
    pool = tuple(iterable)
    n = len(pool)
    if not n and r:
        return
    indices = [0] * r
    yield tuple(pool[i] for i in indices)
    while True:
        for i in reversed(range(r)):
            if indices[i] != n - 1:
                break
        else:
            return
        indices[i:] = [indices[i] + 1] * (r - i)
        yield tuple(pool[i] for i in indices)

combinations_with_replacement() 的代码可被改写为 production() 过滤后的子序列，（相对于元素在输入中的位置）元素不是有序的。

1
2
3
4
5
6
def combinations_with_replacement(iterable, r):
    pool = tuple(iterable)
    n = len(pool)
    for indices in product(range(n), repeat=r):
        if sorted(indices) == list(indices):
            yield tuple(pool[i] for i in indices)

当 n > 0 时，返回项个数为 (n+r-1)! / r! / (n-1)!.

Reference

https://docs.python.org/zh-cn/3.7/library/itertools.html#itertools.islice

-------------本文结束感谢您的阅读-------------