按索引总和对 NumPy 数组的元素进行分组

Question

我有几个尺寸为 30*30*30 的大型 numpy 数组，我需要在其上遍历数组，获取每个索引三元组的总和，并通过该总和将这些元素装箱。 例如，考虑这个简单的 2*2 数组：

test = np.array([[2,3],[0,1]])

该数组具有索引 [0,0]、[0,1]、[1,0] 和 [1,1]。 此例程将返回列表：[2,[3,0],1]，因为数组测试中的2 的索引和为 0，3 和 0 的索引和为 1，1 的索引和为 2。我知道迭代 NumPy 数组并检查总和会起作用，但对于我的实际情况来说，它的效率太低了，有大 N(=30) 和几个数组。 任何有关使用 NumPy 例程完成此分组的输入将不胜感激。 先感谢您。

Answer 1

这是一种应该相当快但不是超快的方法： 30x30x30在我的机器上需要20 ms 。

import numpy as np

# make example
dims = 2,3,4
a = np.arange(np.prod(dims),0,-1).reshape(dims)

# create and sort indices
idx = sum(np.ogrid[tuple(map(slice,dims))])
srt = idx.ravel().argsort(kind='stable')

# use order to arrange and split data
asrt = a.ravel()[srt]
spltpts = idx.ravel().searchsorted(np.arange(1,np.sum(dims)-len(dims)+1),sorter=srt)
out = np.split(asrt,spltpts)

# admire
out
# [array([24]), array([23, 20, 12]), array([22, 19, 16, 11,  8]), array([21, 18, 15, 10,  7,  4]), array([17, 14,  9,  6,  3]), array([13,  5,  2]), array([1])]

Answer 2

您可以程序性地创建索引元组列表并使用它，但可能会进入一个太大而效率低下的代码常量。 [(0,0),[(1,0),(0,1)],(1,1)],

因此，您需要一个函数来为 n 维数组动态生成这些索引。

对于一维，一个简单的计数/增量

   [(0),(1),(2),...] 

第二种，对第一个维度使用一维策略，将第一个维度递减，第二个维度递增填充。

   [(0...)...,(1...)...,(2...)...,...] 
   [[(0,0)],[(1,0),(0,1)],[(2,0),(1,1),(0,2)],[...],...]

请注意，其中一些将在示例数组之外，您的生成器需要包含边界检查。

然后三个维度，对前两个维度进行如上处理，但最后，第一个维度递减，第三个维度递增，重复直到完成

[[(0,0,0),...],[(1,0,0),(0,1,0),...],[(2,0,0),(1,1,0),(0,2,0),...],[...],...]
[[(0,0,0)],[(1,0,0),(0,1,0),(0,0,1)],[(2,0,0),(1,1,0),(0,2,0),(1,0,1),(0,1,1)(0,0,2)

再次需要边界检查或更聪明的起点/终点来避免尝试访问索引之外，但这种通用算法是您如何动态生成索引，而不是让两个大数组竞争缓存和 I/O。

生成 python 或 nympy 等效项留给用户作为练习。