如何获取排序数组的索引，其中前n个值的摘要大于某个值

Question

if I got a sorted numpy array A=[1, 2, 3, 4, 5, 6], I want to get the index where the sum of A[0: index] that bigger than a certain number (supposed to be 9, in this case, the value of index should be 3), how can I write it in python?

I think maybe I can use np.where() or np.cumsum() but it doesn't work

np.argwhere(np.cumsum(A) > 9)

of course, it is wrong, so what is the correct method?

Answer 1

You were close. You can just use argmax as follows. Also, when using argmax, add a check to ensure that atleast 1 value crosses the threshold.

threshold = 9
temp = np.cumsum(A) > threshold

if temp.any():
    print(np.argmax(temp))
else:
    print("Threshold never crossed")

Answer 2

I suggest:

def foo(array, threshold):
    valid_entries = np.cumsum(array) > threshold
    return np.argmax(valid_entries) if np.any(valid_entries) else None