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

我有几个尺寸为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例程来完成此分组的输入都将不胜感激。提前谢谢。

这里有一种方法应该相当快，但不是超快：

30x30x30

在我的机器上需要

20ms

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

---

您可以创建一个索引元组列表并使用它，但可能会进入一个太大而效率低下的代码常量。 [（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等价物留给用户作为练习

相等索引的顺序重要吗？也就是说[2[0,3]，1]可以接受吗？不，子列表中出现的顺序无关紧要。我只需要将它们分组到正确的子列表中。