Cumsum和矢量化切片

时间:2015-04-30 15:05:06

标签: python numpy vectorization

我有一个大小为J的矩阵(j,v)和一个大小为JTildeIDX的向量v。该向量包含 start ,用于J上的cumsum操作。也就是说,给定

>>> JTildeIDX
array([0, 0, 9, 9, 9])
>>> J
array([[ 1.        ,  1.        ,  1.        ,  1.        ,  1.        ],
       [ 1.11111111,  1.11111111,  1.11111111,  1.11111111,  1.11111111],
       [ 1.22222222,  1.22222222,  1.22222222,  1.22222222,  1.22222222],
       [ 1.33333333,  1.33333333,  1.33333333,  1.33333333,  1.33333333],
       [ 1.44444444,  1.44444444,  1.44444444,  1.44444444,  1.44444444],
       [ 1.55555556,  1.55555556,  1.55555556,  1.55555556,  1.55555556],
       [ 1.66666667,  1.66666667,  1.66666667,  1.66666667,  1.66666667],
       [ 1.77777778,  1.77777778,  1.77777778,  1.77777778,  1.77777778],
       [ 1.88888889,  1.88888889,  1.88888889,  1.88888889,  1.88888889],
       [ 2.        ,  2.        ,  2.        ,  2.        ,  2.        ]])

我希望获得大小为v == 5的输出,其中包含以下内容:

np.array([1+1.1+....+2, 1+1.1+....+2, 2, 2, 2])

我将如何实现这一目标?

1 个答案:

答案 0 :(得分:2)

使用braodcastingboolean indexing -

的一种方法
import numpy as np

# Mask of elements from J to be summed up
mask = np.arange(J.shape[0])[:,None] >= JTildeIDX

# Map the mask onto input array and sum along columns for final output
out = (J*mask).sum(0)

示例运行 -

In [154]: J
Out[154]: 
array([[0, 7, 2, 1, 3],
       [0, 8, 4, 3, 1],
       [8, 5, 1, 8, 7],
       [6, 6, 0, 3, 8],
       [6, 5, 6, 4, 2],
       [8, 0, 1, 5, 2],
       [6, 3, 0, 6, 4]])

In [155]: JTildeIDX
Out[155]: array([0, 0, 4, 6, 5])

In [156]: mask = np.arange(J.shape[0])[:,None] >= JTildeIDX

In [157]: mask
Out[157]: 
array([[ True,  True, False, False, False],
       [ True,  True, False, False, False],
       [ True,  True, False, False, False],
       [ True,  True, False, False, False],
       [ True,  True,  True, False, False],
       [ True,  True,  True, False,  True],
       [ True,  True,  True,  True,  True]], dtype=bool)

In [158]: out = (J*mask).sum(0)

In [159]: out
Out[159]: array([34, 34,  7,  6,  6])