python - 按索引之和对 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]，因为数组 test 中的 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])]