def max_indices(arr, k):
'''
Returns the indices of the k first largest elements of arr
(in descending order in values)
'''
assert k <= arr.size, 'k should be smaller or equal to the array size'
arr_ = arr.astype(float) # make a copy of arr
max_idxs = []
for _ in range(k):
max_element = np.max(arr_)
if np.isinf(max_element):
break
else:
idx = np.where(arr_ == max_element)
max_idxs.append(idx)
arr_[idx] = -np.inf
return max_idxs
def largest_indices(ary, n):
"""Returns the n largest indices from a numpy array."""
flat = ary.flatten()
indices = np.argpartition(flat, -n)[-n:]
indices = indices[np.argsort(-flat[indices])]
return np.unravel_index(indices, ary.shape)
15 回答
bottleneck有一个局部排序函数,如果排序整个数组的费用只是为了得到N个最大值太大了 .
我对这个模块一无所知;我只是用Google搜索numpy partial sort .
以下是查看最大元素及其位置的简单方法 . 这里
axis
是域名;axis
= 0表示列方式最大数,axis
= 1表示2D情况的行方式最大数 . 对于更高的尺寸,它取决于你 .较新的NumPy版本(1.8及更高版本)具有一个名为argpartition的功能 . 要获得四个最大元素的索引,请执行
与
argsort
不同,此函数在最坏的情况下以线性时间运行,但返回的索引未排序,从评估a[ind]
的结果可以看出 . 如果您也需要,请在之后对其进行排序:以这种方式按排序顺序获取top-k元素需要O(n k log k)时间 .
我能想到的最简单的是:
这涉及到完整的数组 . 我想知道
numpy
是否提供了一种内置的方式来进行局部排序;到目前为止,我还没有找到一个 .如果这个解决方案变得太慢(特别是对于小型的
n
),那么在Cython中编写一些内容可能是值得的 .我发现使用
np.unique
最直观 .这个想法是,唯一方法返回输入值的索引 . 然后,根据最大唯一值和指标,可以重新创建原始值的位置 .
使用:
对于常规Python列表:
如果您使用Python 2,请使用
xrange
而不是range
.资料来源:heapq — Heap queue algorithm
这将比完整排序更快,具体取决于原始数组的大小和选择的大小:
当然,它涉及篡改原始阵列 . 您可以通过复制或替换原始值来修复(如果需要) . ...以您的用例为准 .
如果你不关心第K个最大元素的顺序,你可以使用argpartition,它应该比通过
argsort
的完整排序表现更好 .积分转至this question .
我运行了一些测试,它看起来像
argpartition
优于argsort
,因为数组的大小和K的值增加 .使用:
它也适用于2D阵列 . 例如,
如果你碰巧使用多维数组,那么你需要展平和解开索引:
例如:
对于多维数组,您可以使用
axis
关键字来沿预期轴应用分区 .并 grab 物品:
但请注意,这不会返回排序结果 . 在这种情况下,您可以沿预期轴使用
np.argsort()
:这是一个例子:
更简单:
其中n是最大值的数量 .
使用:
现在
result
列表将包含 N 元组(index
,value
),其中value
被最大化 .方法
np.argpartition
仅返回k个最大索引,执行本地排序,并且当数组非常大时比np.argsort
(执行完整排序)更快 . 但返回的指数是 NOT in ascending/descending order . 让我们举一个例子:我们可以看到,如果你想要一个严格的升序前k个索引,
np.argpartition
将不会返回你想要的 .除了在np.argpartition之后手动进行排序之外,我的解决方案是使用PyTorch,torch.topk,一种用于神经网络构建的工具,提供类似NumPy的API,同时支持CPU和GPU . 它与使用MKL的NumPy一样快,如果需要大型矩阵/矢量计算,则可以提供GPU提升 .
严格的上升/下降前k个索引代码将是:
请注意torch.topk接受火炬张量,并返回
torch.Tensor
类型中的前k个值和前k个索引 . 与np类似,torch.topk也接受一个axis参数,以便您可以处理多维数组/张量 .我认为最有效的方法是手动迭代数组并保持k大小的最小堆,正如其他人提到的那样 .
而且我也提出了一种蛮力方法:
使用argmax获取其索引后,将最大元素设置为较大的负值 . 接下来argmax的调用将返回第二大元素 . 您可以记录这些元素的原始值并根据需要恢复它们 .