Python 有效地求numpy阵列按索引划分的子阵列的总和

给定一个数组“array”和一组索引“index”，如何找到通过以矢量化方式沿这些索引拆分数组而形成的子数组的累积和？ 为了澄清，假设我有：

>>> 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]])

---

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

您可以使用

np.split

沿索引拆分数组，然后使用python内置函数

map

将

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中的一个内置函数，因此它的性能比常规循环要好。1

以下是二维阵列的备选方案：

>>> 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速度执行的map会使代码以C速度执行。正因为如此，代码是用python实现的，调用函数（第一个参数）并将其应用于iterable项需要时间。

---

您可以在

索引

位置引入原始累积求和数组的微分，以在这些位置创建类似边界的效果，这样，当对微分数组进行累积求和时，将给出停止累积求和的索引输出。乍一看，这可能有点做作，但坚持下去，尝试其他样品，希望会有意义！这个想法与在So中应用的想法非常相似，遵循这样的理念，这里有一种方法与-

这应该是一个几乎矢量化的解决方案，几乎是因为即使我们在循环中计算线性索引，但由于它不是计算密集型部分，因此它对总运行时间的影响将是最小的。此外，如果您不想为了节省内存而破坏输入，则可以将

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]])

---

这太棒了！！！！！这让我感觉好多了。对于2D矩阵，有什么方法可以做到这一点吗？@Aditya369您想如何使用

索引

数组拆分2D矩阵？是要拆分所有行还是？我有一个数组数组，如下所示：index=np.ndarray（[[3,8,14]，[4,9,17]]），对应于一个有两行的数组。因此[3,8,14]代表第一行，[4,9,17]代表第二行。声称

map

“以C速度执行”是非常误导的。

np.hstack（map（np.cumsum，np.split（数组，索引））

和

循环或生成器表达式的常规
，例如
np.hstack（np.cumsum（a）表示np.split（数组，索引））之间的性能基本上没有差异
。即使对于具有100个索引的1000000长数组，使用
map
和生成器表达式在运行时也会有大约13毫秒的差异。这相当于大约14%的速度提升——几乎没有“Python速度”和“C速度”之间的区别。按照您的逻辑，任何只使用CPython内置函数的Python代码都应该“以C速度”运行，这显然是荒谬的。这对2D数组有效吗？索引是否会保持不变？@Aditya369在问题中添加一个2D示例案例和预期输出？我确实添加了这一点。您想让我再添加一个吗？@Aditya369在您实际的2D案例中，2D数组的形状和
索引中的元素数是多少？300*100000是数组的形状，每行的索引数在3000到20000之间，因此总共是900000到6000000个索引。
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]])