python - 从 m by n np.array 中找出 m "smallest"个元素

Question

我有一个大小为 m 的二维 ndarray通过 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 ，以及更好的算法来快速找到这些元素？

Answer 1

我刚刚尝试了另一种方法:

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