Python:按Sum值排列的数组

时间:2017-11-08 10:45:30

标签: python numpy

存在很多主题来根据另一个数组索引对数组求和,但我无法找到我想要的内容。

我有一个数组[1,1,1,1,2,3,3,0,0,2,4],我希望得到的索引到目前为止的总和大于或等于4。所以我会得到:

[3, 5, 9, 10]因为:

Sum([1,1,1,1]) >= 4

Sum([2,3]) >= 4

Sum([3,0,0,2]) >= 4

我需要满足总和值的元素的索引。

我可以通过循环来完成它,但我正在寻找使用numpy等有效地完成它。

2 个答案:

答案 0 :(得分:2)

您可以将np.wherenp.cumsum结合使用:

import numpy as np

li = [1,1,1,1,2,2,3,0,0,1,4]

print(np.where(np.cumsum(li) % 4 == 0))
# (array([ 3,  5,  9, 10], dtype=int32),)

这里的问题当然是它还会找到总和为8,12,16,20等的索引。

不幸的是,我不确定是否有办法“重置”cumsum,如果你有可能做np.where(np.cumsum(li) - 4 == 0)之类的事情(但目前它只会返回3)。

但是,每次cumsum达到4时,您都可以对数组进行切片,但是您不想使用循环。

答案 1 :(得分:1)

array =[1,1,1,1,2,2,3,0,0,1,4]
print(np.where(np.cumsum(array)%4==0))  

现在,如果您更改了数组,那么只要总和超过4就可以检查,那么您可以使用最后一个索引并继续进行。