執行數組計算的Pythonic算法

Question

我正在尋找最有效的pythonic算法來進行數組計算。這裏的問題是：

我有形狀（5,2,3）的陣列，並且其沿着軸線= 0總和如下：

import numpy as np 
A = np.array([[[ 6, 15, 89], 
       [49, 62, 12]], 

       [[92, 8, 34], 
       [93, 81, 35]], 

       [[ 8, 35, 63], 
       [68, 89, 5]], 

       [[27, 20, 85], 
       [87, 42, 90]], 

       [[99, 64, 12], 
       [90, 93, 87]]]) 

B = A.sum(axis=0) 

所以B爲基本上等於A[0]+A[1]+A[2]+A[3]+A[4]它是：

array([[232, 142, 283], 
     [387, 367, 229]]) 

我想知道在求和處理的哪個階段中，各B的6種元素的已經大於100更大的例如元件B[0,0]變得高於100後3個步驟：A[0]+A[1]+A[2]，或B[1,1]後變爲高於100 2步A[0]+A[1]。 所以該算法的最終輸出應該是這樣的數組：

array([[3, 5, 2], 
     [2, 2, 4]]) 

我知道我可以單獨爲每個元素的計算，但我不知道是否有人能想出一個創意和更快的算法。

乾杯，

Answer 1

使用cumsum獲得累積求和，它比較對閾值和最終使用argmax抓住它如同跨越這個門檻的第一個實例 -

(A.cumsum(axis=0) > 100).argmax(axis=0)+1 

採樣運行 -

In [228]: A 
Out[228]: 
array([[[ 6, 15, 89], 
     [49, 62, 12]], 

     [[92, 8, 34], 
     [93, 81, 35]], 

     [[ 8, 35, 63], 
     [68, 89, 5]], 

     [[27, 20, 85], 
     [87, 42, 90]], 

     [[99, 64, 12], 
     [90, 93, 87]]]) 

In [229]: (A.cumsum(0) > 100).argmax(0)+1 
Out[229]: 
array([[3, 5, 2], 
     [2, 2, 4]])