Find column location in matrix based on multiple conditions-CodePudding

I have a matrix in the following form:

import numpy as np
matrix = np.array([[-2,2,6,7,8],[-3,7,1,0,-2]])

I want to find the location of the column with the highest possible value in the first row conditional on non-negative numbers in the second row e.g. in my case I want the algorithm to find the 4th row.

solution = np.array([7,0])
column_location = 3

I tried using numpy functions like np.min(), np.max(),np.take() but I loose the location information when subsampling the matrix.

CodePudding user response：

Simply:

nn = np.where(matrix[1] >= 0)[0]
ix = nn[matrix[0, nn].argmax()]

On your data:

>>> ix
3

CodePudding user response：

Here's a sketch:

pos_inds = np.where(matrix[1, :] >= 0)[0]  # indices where 2nd row is positive

max_ind = matrix[0, pos_inds].argmax()  # max index into row with positive values only

orig_max_ind = pos_inds[max_ind]  # max index into the original array


print(orig_max_ind)  # 3
print(matrix[:, orig_max_ind])  # [7, 0]

CodePudding user response：

Here I use the masks to handle the numpy, and also consider that if all the numbers are negative in second column, there will be no solution:

import numpy as np
import numpy.ma as ma
from copy import deepcopy

min_int = -2147483648

matrix = np.array([[-2, 2, 6, 7, 8], [-3, 7, 1, 0, -2]])

# we keep the original matrix untouched
matrix_copy = deepcopy(matrix)
masked_array = ma.masked_less(matrix[1], 0)
matrix_copy[0][masked_array.mask] = min_int
column_location = np.argmax(matrix_copy[0])

if matrix_copy[0][column_location] == min_int:
    print("No solution")
else:
    solution = np.array([matrix[0][column_location], matrix[1][column_location]])
    print(solution)  # np.array([7,0])
    print(column_location)  # 3