python - 快速查找二维数组中的多个最大值

Question

情况如下:

我有一个二维 numpy 数组。它的形状是(1002, 1004)。每个元素包含一个介于 0 和 Inf 之间的值。我现在要做的是确定前 1000 个最大值并将相应的索引存储到名为 x 的列表和名为 y 的列表中。这是因为我想绘制最大值，而索引实际上对应于值的实时 x 和 y 位置。

我目前拥有的是:

x = numpy.zeros(500)
y = numpy.zeros(500)

for idx in range(500):
    x[idx] = numpy.unravel_index(full.argmax(), full.shape)[0]
    y[idx] = numpy.unravel_index(full.argmax(), full.shape)[1]
    full[full == full.max()] = 0.

print os.times()

这是我的 2D numpy 数组。从for循环可以看出，我暂时只确定前500个最大值。然而，这已经花费了大约 5 秒。对于前 1000 个最大值，用户时间实际上应该在 0.5 秒左右。我注意到一个非常耗时的部分是每次都将先前的最大值设置为 0。我怎样才能加快速度？

非常感谢!

Answer 1

如果你有 numpy 1.8，你可以使用 argpartition函数或方法。 下面是计算 x 和 y 的脚本:

import numpy as np

# Create an array to work with.
np.random.seed(123)
full = np.random.randint(1, 99, size=(8, 8))

# Get the indices for the largest `num_largest` values.
num_largest = 8

indices = (-full).argpartition(num_largest, axis=None)[:num_largest]
# OR, if you want to avoid the temporary array created by `-full`:
# indices = full.argpartition(full.size - num_largest, axis=None)[-num_largest:]

x, y = np.unravel_index(indices, full.shape)

print("full:")
print(full)
print("x =", x)
print("y =", y)
print("Largest values:", full[x, y])
print("Compare to:    ", np.sort(full, axis=None)[-num_largest:])

输出:

full:
[[67 93 18 84 58 87 98 97]
 [48 74 33 47 97 26 84 79]
 [37 97 81 69 50 56 68  3]
 [85 40 67 85 48 62 49  8]
 [93 53 98 86 95 28 35 98]
 [77 41  4 70 65 76 35 59]
 [11 23 78 19 16 28 31 53]
 [71 27 81  7 15 76 55 72]]
x = [0 2 4 4 0 1 4 0]
y = [6 1 7 2 7 4 4 1]
Largest values: [98 97 98 98 97 97 95 93]
Compare to:     [93 95 97 97 97 98 98 98]