Finding the index of max value of columns in numpy array but removing the previous max-CodePudding

I have an array with N rows and M columns.

I would like to run through all the columns, finding the index of the row in which contains the max value of the column. However, each row should be selected only once.

For instance, let's consider a matrix

1 1
2 2

The output should be [1, 0]. Because the row 1 (value of 2) is the max value of column 0, then we move to column 2, the row 1 is out of consideration, so the row 0 will be the highest cell.

Indeed, things can be solved easily with for a nested for loop, and something like:

removed_rows = []
for i in range (nb_columns):
  index_max = 0
  value_max = A[0,i]
  for j in range (nb_rows):
    if j in removed_rows:
      continue
    else:
      if value_max < A[j,i]:
        index_max = j
        value_max = A[j,i]
  removed_rows.append (index_max)

However, it seems slow for a huge matrix. Is there any method we can do it faster (with numpy?)?

Many thanks

CodePudding user response：

This might not be very fast as it still loop through the columns, which I think is unavoidable due to the constrain, but should be faster than your solution as it finds the maximum's index with argmax:

out = []
mm = A.min() - 1

for j in range(A.shape[1]):
    idx = np.argmax(A[:,j])

    # replace the entire row with mm
    # so next `argmax` will ignore this row
    A[idx] = mm
    out.append(idx)

The above takes about 640 us on 100 x 100 arrays, and 18ms on 1k x 1k arrays. Your code refuses to run on 1k x 1k array within reasonable time on my system.