我有一个大小为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
,以及更好的算法来快速找到这些元素?
答案 0 :(得分: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