从n个np.array中找到m个“最小”元素

Rain Lee 发表于 Dev

李雨

我有一个二维数组ndarray，大小m为n（m<=n），如下所示：

a = [ [1, 2, 3], 
      [4, 5, 6] ]

现在，我想m从数组中贪婪地找到“最小”元素，并限制每个行和列只能选择一个元素，每次选择全局最小值。我的代码如下：

for k in xrange(m):
    index = np.argmin(a)
    i, j = divmod(index, n-k)
    result.append(a[i][j])
    a = np.delete(np.delete(a, i, 0), j, 1)

所以我会得到result = [1, 5]，有没有更好的方法来表示输入数组a，以及找到速度更快的这些元素的更好算法？

彼得·韦普雷克

我只是尝试了另一种方法：

import numpy as np
import timeit

nmin = 2000 # number of the smallest values to find in a matrix with unique row and column indexes
nrows = 2000 # number of rows
ncols = 2000 # number of columns
print "Select {} smallest values from {} x {} matrix".format(nmin, nrows, ncols)

matrix = np.random.uniform(0, 1, size = nrows * ncols).reshape(nrows, ncols) # sample 2D array
#print matrix

# ALTERNATIVE: sort once and track-and-skip visited rows and columns
startedat = timeit.default_timer()
seenrows = set()
seencols = set()
order = (divmod(index, ncols) for index in np.argsort(matrix, None))
for iter in xrange(nmin):
    while True:
        try:
            current = order.next()
        except:
            break
        if current[0] not in seenrows and current[1] not in seencols:
            #print iter, current, matrix[current[0]][current[1]]
            seenrows.add(current[0])
            seencols.add(current[1])
            break
alternative = timeit.default_timer() - startedat
print "Alternative approach took: ", alternative

# ORIGINAL: repeatedly find minimum and update matrix
startedat = timeit.default_timer()
for k in xrange(nmin):
    index = np.argmin(matrix)
    i, j = divmod(index, np.shape(matrix)[1])
    #print k, (i, j), matrix[i][j]
    matrix = np.delete(np.delete(matrix, i, 0), j, 1)
    if matrix.size == 0: break
original = timeit.default_timer() - startedat
print "   Original approach took: ", original, "WINNER" if original < alternative else "TIE" if original == alternative else "LOOSER"

结果如下：

Select 2 smallest values from 2000 x 2000 matrix
Alternative approach took:  0.737312265981
   Original approach took:  0.0572765855289 WINNER

Select 20 smallest values from 2000 x 2000 matrix
Alternative approach took:  0.732718787079
   Original approach took:  0.564769882057 WINNER

Select 200 smallest values from 2000 x 2000 matrix
Alternative approach took:  0.736015078962
   Original approach took:  5.14679721535 LOOSER

Select 2000 smallest values from 2000 x 2000 matrix
Alternative approach took:  6.46196502191
   Original approach took:  19.2116744154 LOOSER

Select 20000 smallest values from 2000 x 2000 matrix
Alternative approach took:  7.90157398272
   Original approach took:  19.189003763 LOOSE