python - 在numpy矩阵中查找具有不同列索引的行式最大值列

Question

我有一个大 numpy 10,000 行和 10,000 列的矩阵 float类型。对于每一行，我需要找到具有最大值的列索引并最多选择每列一次。
例如，对于给定的数组 arr ，我需要输出为 list/array的 (row index, column index) out 中给出的元组:

arr = np.array(
   [[0.86, 0.23, 0.83, 0.79],
   [0.15, 0.98, 0.86, 0.47],
   [1.  , 0.08, 0.01, 0.04],
   [0.78, 0.82, 0.17, 0.56],
   [0.73, 0.91, 0.52, 0.31]])
out = [(0,0),(1,1),(2,3),(3,2)]

解释:

最初 out是空的。

对于 row 0 ，最大值为 0.86在 column 0 ，所以 out现在是 [(0,0)]

对于 row 1 ，最大值为 0.98在 column 1 , 和 column 1还没有出现，所以 out现在是 [(0,0),(1,1)]

对于 row 2 ，最大值为 1在 column 0 ，但是 column 0已经被选中，所以我们寻找下一个最大值，即 0.08在 column 1 ，也存在于 out 中，然后是下一个最大值，即 0.04在 column 3 , 所以 out现在是 [(0,0),(1,1),(2,3)]

同样，对于 row 3 , 尚未选择的最大值列是 column 2 ，所以最后 out是 [(0,0),(1,1),(2,3),(3,2)]

我想尽可能有效地计算它。 O(n 2 )使用 2 for 的解决方案循环是微不足道的，所以任何比这更好的解决方案(使用内置 numpy 函数更好的时间复杂度或更好的运行时间)都会非常有帮助。

Answer 1

如果您愿意使用 cython，则可以稍微加快迭代速度。

# distutils: language = c++
# cython: boundscheck = False

from libcpp.set cimport set as cset
from libc.math cimport INFINITY

def method2(double[:, :] x):
    cdef:
        int nrows, ncols, i, j, best_j
        double best_value
        # Define the set of columns that have already been used.
        cset[int] usedcols
        
    nrows = x.shape[0]
    ncols = x.shape[1]
    out = []
    
    for i in range(nrows):
        best_value = -INFINITY
        best_idx = -1
        
        # Find the largest value for each row that's not already used.
        for j in range(ncols):
            if x[i, j] > best_value and usedcols.find(j) == usedcols.end():
                best_value = x[i, j]
                best_j = j
        out.append((i, best_j))
        usedcols.insert(best_j)
    
    return out

假设 Samarth 的解决方案是 method1 ，这是一个性能比较。

x = np.random.normal(0, 1, (10000, 10000))
%timeit method1(x)  # 770 ms ± 4.98 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit method2(x)  # 57.7 ms ± 149 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)