Finding row-wise maximum value column with distinct column indices in a numpy matrix

user_8035804 0
Viewed 346

I have a large numpy matrix of the order of 10,000 rows and 10,000 columns of float type. For each row, I need to find the column index with maximum value and select each column at most once.

For example, for the given array arr, I need the output as a list/array of (row index, column index) tuples as given in 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)]

Explanation:

  • Initially out is empty.
  • For row 0, the maximum value is 0.86 in column 0, so out is now [(0,0)]
  • For row 1, the maximum value is 0.98 in column 1, and column 1 isn't already present in out, so out is now [(0,0),(1,1)]
  • For row 2, the maximum value is 1 in column 0, but column 0 has already been selected, so we look for next maximum, i.e. 0.08 in column 1, which is also present in out, then next maximum, i.e. 0.04 in column 3, so out is now [(0,0),(1,1),(2,3)]
  • Similarly, for row 3, the column with maximum value that hasn't been already selected is column 2, so final out is [(0,0),(1,1),(2,3),(3,2)]

I want to compute it as efficiently as possible. The O(n2) solution using 2 for loops is trivial, so any solution better than this (either better time complexity or better runtime using inbuilt numpy functions) would be very helpful.

python-3.x numpy arrays python performance
2 Answers

If you're open to using cython, the iteration can be sped up somewhat.

# 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

Assuming Samarth's solution is method1, here's a performance comparison.

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)
user_11509610

Assuming the minimum value in the matrix is 0

for i in range(len(arr)): # for the ith row
    maxcol = np.argmax(arr[i])
    arr[:,maxcol] = np.zeros( len(arr[:,maxcol]) )
    out.append((i,maxcol))

This works by finding the index of the max column for a row (using argmax), and then setting all of that column's elements to an arbitrary minimum value (0 here, given efficiently by np.zeros)

This should work, I'm trying it out rn. And is definitely less than O(n^2) but don't know exactly how much.

Edit: I tried it and out is [(0, 0), (1, 1), (2, 3), (3, 2), (4, 0)]. So it does have the desired output, but an extra term (4,0). [which is there because iterates over all rows]. If you want to iterate till you 'run out' of columns, you can use min(arr.shape) instead of len(arr)

Time Complexity

We're running a for loop once, that takes O(n)
Inside the for loop, we have np.argmax, which is also O(n)
Then comes the replacements which are constant time [O(1)]

So it's not that optimised. You can write your own code for that, by directly coding the edge cases and other improvements like checking only the required parts, etc...

user_132643340
Related