查找由numpy數組的索引分割的子數組的子集

Question

給定一個數組'數組'和一組索引'索引'，如何找到通過以矢量化方式沿着這些索引分割數組而形成的子數組的累積和？ 澄清一下，假設我有：

>>> array = np.arange(20)
>>> array
array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19])
indices = np.arrray([3, 8, 14])

操作應該輸出：

array([0, 1, 3, 3, 7, 12, 18, 25, 8, 17, 27, 38, 50, 63, 14, 29, 45, 62, 80, 99])

請注意，陣列非常大（100000個元素），因此，我需要一個矢量化答案。 使用任何循環會大大減慢它。 另外，如果我有同樣的問題，但是2D數組和相應的索引，我需要為數組中的每一行做同樣的事情，我該怎么辦呢？

對於2D版本：

>>>array = np.arange(12).reshape((3,4))
>>>array
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
>>> indices = np.array([[2], [1, 3], [1, 2]])

輸出將是：

array([[ 0,  1,  3,  3],
       [ 4,  9,  6, 13],
       [ 8, 17, 10, 11]])

澄清：每一行都將被拆分。

Answer 1

您可以在indices位置引入原始累積求和數組的微分，以在這些位置創建類似邊界的效果，這樣當差分數組被累加求和時，為我們提供索引停止的累積求和輸出。 這可能在初看起來有點做作 ，但堅持下去，試試其他樣品，希望有意義！ 這個想法非常類似於this other MATLAB solution.應用的想法this other MATLAB solution. 所以，遵循這樣一種理念，這里使用numpy.diff和cumulative summation的一種方法 -

# Get linear indices
n = array.shape[1]
lidx = np.hstack(([id*n+np.array(item) for id,item in enumerate(indices)]))

# Get successive differentiations
diffs = array.cumsum(1).ravel()[lidx] - array.ravel()[lidx]

# Get previous group's offsetted summations for each row at all 
# indices positions across the entire 2D array
_,idx = np.unique(lidx/n,return_index=True)
offsetted_diffs = np.diff(np.append(0,diffs))
offsetted_diffs[idx] = diffs[idx]

# Get a copy of input array and place previous group's offsetted summations 
# at indices. Then, do cumulative sum which will create a boundary like 
# effect with those offsets at indices positions.
arrayc = array.copy()
arrayc.ravel()[lidx] -= offsetted_diffs
out = arrayc.cumsum(1)

這應該是一個幾乎矢量化的解決方案，幾乎是因為即使我們在循環中計算線性索引，但由於它不是這里的計算密集部分，因此它對總運行時間的影響將是最小的。 此外，如果您不關心破壞輸入以保存內存，則可以將arrayc替換為array 。

樣本輸入，輸出 -

In [75]: array
Out[75]: 
array([[ 0,  1,  2,  3,  4,  5,  6,  7],
       [ 8,  9, 10, 11, 12, 13, 14, 15],
       [16, 17, 18, 19, 20, 21, 22, 23]])

In [76]: indices
Out[76]: array([[3, 6], [4, 7], [5]], dtype=object)

In [77]: out
Out[77]: 
array([[ 0,  1,  3,  3,  7, 12,  6, 13],
       [ 8, 17, 27, 38, 12, 25, 39, 15],
       [16, 33, 51, 70, 90, 21, 43, 66]])

Answer 2

您可以使用np.split沿索引拆分數組，然后使用內置函數map python將np.cumsum()應用於子數組。 最后通過使用np.hstack將結果轉換為集成數組：

>>> np.hstack(map(np.cumsum,np.split(array,indices)))
array([ 0,  1,  3,  3,  7, 12, 18, 25,  8, 17, 27, 38, 50, 63, 14, 29, 45,
       62, 80, 99])

請注意 ，由於map是python中的內置函數，並且已經在Python解釋器中的C中實現，因此它的性能優於常規循環。 ¹

以下是2D陣列的替代方案：

>>> def func(array,indices):
...     return np.hstack(map(np.cumsum,np.split(array,indices)))
... 
>>> 
>>> array
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
>>> 
>>> indices
array([[2], [1, 3], [1, 2]], dtype=object)

>>> np.array([func(arr,ind) for arr,ind in np.array((array,indices)).T])
array([[ 0,  1,  2,  5],
       [ 4,  5, 11,  7],
       [ 8,  9, 10, 21]])

請注意 ，您的預期輸出不是基於np.split工作方式。

如果您想要這樣的結果，您需要在索引中添加1：

>>> indices = np.array([[3], [2, 4], [2, 3]], dtype=object)
>>> 
>>> np.array([func(arr,ind) for arr,ind in np.array((array,indices)).T])
array([[  0.,   1.,   3.,   3.],
       [  4.,   9.,   6.,  13.],
       [  8.,  17.,  10.,  11.]])

由於評論說使用生成器表達式和map函數之間沒有性能差異，我運行了一個更好地演示結果的基准測試。

# Use map
~$ python -m timeit --setup "import numpy as np;array = np.arange(20);indices = np.array([3, 8, 14])" "np.hstack(map(np.cumsum,np.split(array,indices)))"
10000 loops, best of 3: 72.1 usec per loop
# Use generator expression
~$ python -m timeit --setup "import numpy as np;array = np.arange(20);indices = np.array([3, 8, 14])" "np.hstack(np.cumsum(a) for a in np.split(array,indices))"
10000 loops, best of 3: 81.2 usec per loop

---

_{請注意，這並不意味着使用以C速度執行的映射會使該代碼以C速度執行。 這是因為，代碼已經在python中實現並調用函數（第一個參數）並將其應用於可迭代項目需要時間。}