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

【问题标题】：Grouping elements of a NumPy array by sum of indices按索引总和对 NumPy 数组的元素进行分组
【发布时间】：2020-07-03 03:47:45
【问题描述】：

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

【问题讨论】：

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

标签： python arrays numpy

【解决方案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])]

【讨论】：

【解决方案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)

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

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

【讨论】：