python - itertools.groupby 中组大小的限制

Question

我正在寻找一种解决方案，如何限制由 itertools.groupby 创建的组的组大小。

目前我有这样的东西:

>>> s = '555'
>>> grouped = groupby(s)
>>> print([(k, len(list(g))) for k, g in grouped])
[('5', 3)]

我想要实现的是最大组的大小=2，所以我的输出将是:

[('5', 2), ('5', 1)]

有没有简单有效的方法来做到这一点？也许以某种方式通过提供给 groupby 的 key 参数？

Answer 1

这是一个使用 groupby 和 defaultdict 的解决方案。

from itertools import groupby
from collections import defaultdict

s = "5555444"
desired_length = 2
counts = defaultdict(int)

def count(x):
    global counts
    c = counts[x]
    counts[x] += 1
    return c

grouped = groupby(s, key=lambda x: (x, count(x) // desired_length))
print([(k[0], len(list(g))) for k, g in grouped])

老实说，我认为这种解决方案是 Not Acceptable ，因为它要求您始终跟踪全局状态，但它就是这样。我个人只会使用类似缓冲区的东西。

from collections import defaultdict
s = "5555444"

def my_buffer_function(sequence, desired_length):
    buffer = defaultdict(int)
    for item in sequence:
        buffer[item] += 1
        if buffer[item] == desired_length:
            yield (item, buffer.pop(item))
    for k, v in buffer.items():
        yield k, v

print(list(my_buffer_function(s, 2)))

这也是一个生成器。但它可能会遗漏一些你目前依赖的 groupby 东西。