I have a numpy array of shape (9, 200, 200). I would like to get a list of the indices of the minimum value for each of the 0th dimension rows. For example, my output would be a 200 by 200 array with each element being a list of the indices of the minimum value for the 0th dimension row ([0, 2, 3] etc.). I need the solution to work with multiple minimum values per row so that I have a list of multiple indices per row. Time is a factor, so I would prefer a solution that is faster than basic python for loops.
zero_array = np.random.rand(9, 200, 200) min_row_values = np.where(zero_array == np.min(zero_array))
My current solution is posted above, but this code just shows the minimum index value of the entire array rather than the minimum indices of each row.
Advertisement
Answer
argmin will give you the index of the first minimum element in each column. Numpy does not generally support ragged arrays, so you will need to jump through extra hoops to get the lists you want.
A good way to start is probably to create a mask of the elements of interest:
mask = (x == x.min(axis=0, keepdims=True))
You can convert the mask to indices using something like np.nonzero. I would probably do it on a partially raveled array:
row, col = np.nonzero(mask.reshape(mask.shape[0], -1))
You can then take row and split it on unique values of col, so that you have arrays of indices in each column, arranged in a list. You will need to sort on col first though:
sort_order = col.argsort() row = row[sort_order] col = col[sort_order]
A more clever (but less general) approach, is to call np.nonzero with the dimensions reversed so that the automatic sort on the first index takes care of all that for you:
col, row = np.nonzero(mask.reshape(mask.shape[0], -1).T)
The actual split is pretty easy now:
row = np.split(row, np.flatnonzero(np.diff(col)) + 1)
You can then emplace the resulting indices into an array of dtype=object to allow raggedness:
result = np.reshape(row, x.shape[1:])
This simplified approach will raise a warning in more recent versions of numpy, so you should probably do the following instead:
result = np.empty(x.shape[1:], dtype=object) result.ravel()[:] = row