Python如何依据另一可迭代对象长度高效切片可迭代对象

核心思路就是利用zip配对两个迭代器的特性，刚好取到和分组g等长的values元素，修正后的写法非常简洁：

from itertools import groupby

x = [(5, 65), (2, 12), (5, 18), (3, 35), (4, 49), (4, 10), (1, 27), (1, 1), (4, 71), (2, 41), (2, 17), (1, 25), (2, 62), (5, 65), (4, 5), (1, 51), (1, 13), (5, 92), (2, 62), (5, 81)]

keys, values = map(iter, zip(*sorted(x)))
print([sum(v for v, _ in zip(values, g)) for _, g in groupby(keys)])
# 输出: [117, 194, 35, 135, 321]

原理解释 #

• zip(values, g)会同时遍历两个迭代器，每次返回一个(values当前元素, g当前元素)的元组，直到任意一个迭代器耗尽为止。因为g的长度就是当前分组的元素数，所以刚好会从values中取对应数量的元素，不会多消耗也不会少取。
• 此前尝试的next(zip(*zip(values, g)))写法错误，原因是每次只取了zip生成的第一个元组的第一个元素，剩下的配对元素被直接丢弃，导致多消耗了values的元素，最终结果不符合预期。

函数式风格变体 #

如果偏好纯函数式写法，可以结合map和operator.itemgetter实现，效率和列表推导式基本一致：

from operator import itemgetter
print([sum(map(itemgetter(0), zip(values, g))) for _, g in groupby(keys)])

其他方案说明 #

用islice的方案需要先把g转成列表获取长度，多了一次遍历分组的开销，反而不如zip方案高效简洁。另外上述zip配对的方案可以直接复用到其他「按另一个可迭代对象长度截取当前迭代器元素」的场景，泛用性很强。

内容的提问来源于stack exchange，提问作者Jab