我有一个ndarray。从这个数组中,我需要选择具有最大值的 N 个数字的列表。我发现heapq.nlargest
找到了 N 个最大的条目,但我需要提取索引。我想构建一个新数组,其中只有第一列中权重最大的 N 行才能存活。其余行将被随机值替换
import numpy as np
import heapq # For choosing list of max values
a = [[1.1,2.1,3.1], [2.1,3.1,4.1], [5.1,0.1,7.1],[0.1,1.1,1.1],[4.1,3.1,9.1]]
a = np.asarray(a)
maxVal = heapq.nlargest(2,a[:,0])
if __name__ == '__main__':
print a
print maxVal
我的输出是:
[[ 1.1 2.1 3.1]
[ 2.1 3.1 4.1]
[ 5.1 0.1 7.1]
[ 0.1 1.1 1.1]
[ 4.1 3.1 9.1]]
[5.0999999999999996, 4.0999999999999996]
但我需要的是[2,4]
作为构建新数组的索引。索引是行,所以如果在这个例子中我想用 0 替换其余部分,我需要完成:
[[0.0 0.0 0.0]
[ 0.0 0.0 0.0]
[ 5.1 0.1 7.1]
[ 0.0 0.0 0.0]
[ 4.1 3.1 9.1]]
我被困在需要索引的地方。原始数组有 1000 行和 100 列。权重是归一化的浮点数,我不想做类似的事情,if a[:,1] == maxVal[0]:
因为有时我的权重非常接近,并且可以完成maxVal[0]
比我原来的 N更多的值。
是否有任何简单的方法可以在此设置中提取索引以替换数组的其余部分?