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

时间:2013-12-29 14:13:41

标签: python arrays performance numpy max

情况如下:

我有一个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。我怎样才能加快速度?

非常感谢你!

4 个答案:

答案 0 :(得分:12)

如果你有numpy 1.8,你可以使用argpartition功能或方法。 这是一个计算xy的脚本:

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]

答案 1 :(得分:2)

你可以像@Inspired建议的那样遍历数组,但是逐项循环遍历NumPy数组会导致代码比使用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

显示这增加了大约4毫秒到nlargest_indices_orig的5s基准。


警告:如果nlargest_indices包含重复值,则nlargest_indices_origarr可能会返回不同的结果。

nlargest_indicesn中找到arr个最大值,然后返回与这些值的位置对应的xy索引。

nlargest_indices_orign中找到arr最大值,然后为每个大值返回一个 xy索引。如果有多个xy对应相同的大值,则可能会遗漏某些值较大的位置。

他们也以不同的顺序返回索引,但我认为这与你的绘图目的无关。

答案 2 :(得分:1)

如果你想知道2d数组中n max / min值的索引,我的解决方案(最大的是)

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

这将找到来自展平数组的最大值的索引,然后根据余数和mod确定2d数组中的索引。

没关系。基准测试表明,至少对于num_largest = 3,解开方法的速度是其两倍。

答案 3 :(得分:-1)

我担心最耗时的部分是重新计算最大值。事实上,您必须计算最多1002 * 1004个数字500次,这样可以进行5亿次比较。

可能你应该编写自己的算法来一次性找到解决方案:在扫描你的2D数组时,只保留1000个最大数字(或它们的索引)(不修改源数组)。我认为某种二进制堆(看看heapq)适合存储。