在2d数组中快速找到多个最大值

Question

情况如下：

我有一个2D的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。如何加快速度呢？

非常感谢！

- The Dude

---

你能使用numpy 1.8吗？它有partition和argpartition函数。我在我的答案中使用了后者。 - Warren Weckesser

谢谢您的回答，但我不能使用numpy 1.8。partition和unique相似吗？ - The Dude

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]

- Warren Weckesser

---

+1 这可能是 1.8 中最酷的算法之一。虽然从示例中很明显，但在 OP 的问题中并不是必需的，但值得强调的是 partition 函数不会将分割数组的块排序。 - Jaime

Answer 2

您可以像@Inspired建议的那样遍历数组，但是逐个遍历NumPy数组往往会导致性能较慢的代码，因为NumPy函数是用C / Fortran编写的，而逐个遍历则倾向于使用Python函数。

因此，尽管排序的时间复杂度为O(n log n)，但它可能比基于Python的一次遍历O(n)的解决方案更快。下面的np.unique执行排序：

import numpy as np

def nlargest_indices(arr, n):
    uniques = np.unique(arr)
    threshold = uniques[-n]
    return np.where(arr >= threshold)

full = np.random.random((1002,1004))
x, y = nlargest_indices(full, 10)
print(full[x, y])
print(x)
# [  2   7 217 267 299 683 775 825 853]
print(y)
# [645 621 132 242 556 439 621 884 367]

这里是一个时间基准测试，比较了nlargest_indices（上面）和以下内容：

def nlargest_indices_orig(full, n):
    full = full.copy()
    x = np.zeros(n)
    y = np.zeros(n)

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


In [97]: %timeit nlargest_indices_orig(full, 500)
1 loops, best of 3: 5 s per loop

In [98]: %timeit nlargest_indices(full, 500)
10 loops, best of 3: 133 ms per loop

为了进行时间测量，我需要复制nlargest_indices_orig中的数组，以免full在计时循环中被改变。

对复制操作进行基准测试：

def base(full, n):
    full = full.copy()

In [102]: %timeit base(full, 500)
100 loops, best of 3: 4.11 ms per loop

这个测试显示，使用nlargest_indices_orig大约会使5秒的基准测试增加4毫秒。

---

注意：如果arr包含重复值，nlargest_indices和nlargest_indices_orig可能会返回不同的结果。

nlargest_indices查找arr中前n个最大的值，并返回对应于这些值位置的x和y索引。

nlargest_indices_orig查找arr中前n个最大的值，并为每个大值返回一个 x 和 y 的索引。如果有多个x 和 y 对应于相同的大值，则可能会错过一些出现大值的位置。

它们还以不同的顺序返回索引，但我想这对您绘图的目的没有影响。

- unutbu

---

这个方法很有效。非常感谢！我的二维数组中确实有相同的最大值，但是还没有遇到任何问题。 - The Dude

Answer 3

如果您想知道二维数组中 n 个最大/最小值的索引，我的解决方案（对于最大值）是：

indx = divmod((-full).argpartition(num_largest,axis=None)[:3],full.shape[0])

这段代码会在扁平化的数组中找到最大值的索引，然后根据余数和模数确定2D数组中的索引位置。

算了，基准测试显示unravel方法至少比num_largest = 3时快两倍。

- Tom Logan

Answer 4

我担心最耗时间的部分是重新计算最大值。实际上，你需要计算1002 * 1004个数的最大值500次，这将给出5亿次比较。

也许你应该编写自己的算法一次性找到解决方案：在扫描二维数组时（不修改源数组），将只保留1000个最大数字（或其索引）放在某个地方。我认为某种类型的二叉堆（看看heapq)适合用于存储。

- nullptr