[FIXED] numpy argmax when values are equal

Issue

I got a numpy matrix and I want to get the index of the maximum value in each row. E.g.

[[1,2,3],[1,3,2],[3,2,1]]

will return

[0,1,2]

However, when there're more than 1 maximum value in each row, numpy.argmax will only return the smallest index. E.g.

[[0,0,0],[0,0,0],[0,0,0]]

will return

[0,0,0]

Can I change the default (smallest index) to some other values? E.g. when there're equal maximum values, return 1 or None, so that the above result will be

[1,1,1]
or
[None, None, None]

If I can do this in TensorFlow that'll be better.

Thanks!

Solution

A convenient value of "default" is -1, as argmax will not return that on its own. None does not fit in an integer array. A masked array is also an option, but I didn't go that far. Here is a NumPy implementation

def my_argmax(a):
    rows = np.where(a == a.max(axis=1)[:, None])[0]
    rows_multiple_max = rows[:-1][rows[:-1] == rows[1:]]
    my_argmax = a.argmax(axis=1)
    my_argmax[rows_multiple_max] = -1
    return my_argmax

Example of use:

import numpy as np
a = np.array([[0, 0, 0], [4, 5, 3], [3, 4, 4], [6, 2, 1]])
my_argmax(a)   #  array([-1,  1, -1,  0])

Explanation: where selects the indexes of all maximal elements in each row. If a row has multiple maxima, the row number will appear more than once in rows array. Since this array is already sorted, such repetition is detected by comparing consecutive elements. This identifies the rows with multiple maxima, after which they are masked in the output of NumPy's argmax method.

Answered By - user6655984