python - 如何编写一个函数来仅计算类数组对象的前 n 项并返回它们的和？

Question

假设我有一个如下数组:

A = np.array([a,b,c,d,e,f])

数组中的每个元素都有一个等于标量的值(例如，a 可能等于 2，b 等于 7，等等)。但是 a 和 b 需要很长时间来计算，因此 Python 可能需要几分钟才能确定值 a，更多的时间来确定 b 的值，对于其余元素也是如此。

是否有办法定义一个函数(我们称之为 sumfirstn)来计算 A 的前 n 项并计算它们的总和？例如，sumfirstn(A,2) 将返回 a+b，sumfirstn(A,3) 将返回 a+b +c 等

但是(这一点很重要)sumfirstn(A,2) 必须不评估c、d 、e 和 f。它必须仅评估正在添加的相关项(以减少计算时间)。当然，A[0:n].sum() 可以解决这个问题，但这需要 Python 计算整个数组 A，而我不想这样做!

有什么想法吗？ A 不一定是数组，可以是列表、数据框或其他某种形式......无论什么都行。

Answer 1

看起来您需要一个通用的解决方案来对任何可迭代的前n项进行求和，它不一定必须支持索引、切片或求和 方法。

如果我理解正确，以下应该足以作为通用解决方案。

>>> from itertools import islice
>>> 
>>> def first_n_sum(iterable, n):
...     return sum(islice(iterable, n))

演示:

>>> first_n_sum([1,2,3,4], 3)
6
>>> first_n_sum((x for x in (1, 2, 3, 4)), 2)
3
>>> first_n_sum([], 123)
0
>>> from collections import OrderedDict
>>> first_n_sum(OrderedDict([(1, 'a'), (2, 'b'), (-3, 'c'), (100, 'd')]), 3)
0

请注意，如果可迭代对象无法生成 n 个值，则不会出现错误，就像切片超出数组长度是可以的一样。在这种情况下，将返回所有元素的总和。