Is there a way to find the UNIQUE row indices of maximum columnar values in a 2D NumPy array?

user_2269232 0
Viewed 183

For each column in a 2D NumPy array, the column's maximum value can appear more than once. I would like to find the row index for each column maximum, without repeating row indices.

Here is an example that demonstrates why np.argmax doesn't work:

import numpy as np

a = np.array([[1, 1, 0],
              [1, 0, 1],
              [0, 0, 1]])

ind = np.argmax(a, axis=0)

print(ind)

Output:

[0 0 2]

I want the result: [1, 0, 2] for this example.

That is:

  • The row index for the second column must be 0
  • This implies that the row index for the first column must be 1
  • This in turn implies that the row index for the third column must be 2

A slightly more complex example is this array:

a = np.array([[1, 1, 0],
              [1, 1, 1],
              [0, 0, 1]])

In this case, there is no column with a unique maximum value. I'd be happy with either of these answers:

  • [0, 1, 2]
  • [1, 0, 2]

An even more complex example is:

a = np.array([[1, 1, 1],
              [1, 1, 1],
              [0, 1, 1]])

In this case, I'd be happy with any of these answers:

  • [0, 1, 2]
  • [0, 2, 1]
  • [1, 0, 2]
  • [1, 2, 0]

I can solve these problems with loops and logical conditions, but I'm wondering if there is a way to solve the problem using numpy functions?

argmax multidimensional-array python-3.x numpy arrays
2 Answers

May be overkill, but you can use scipy.optimize.linear_sum_assignment:

from scipy.optimize import linear_sum_assignment

a = np.array([[1, 1, 0],
              [1, 0, 1],
              [0, 0, 1]])

linear_sum_assignment(-a.T)[1]
# array([1, 0, 2])

Note that you can always reduce to the 0,1 case using something like

abin = a==a.max(axis=0)

This can speed up the assignment quite a bit.

Alternatively, see this post for a graph theory solution.

user_72073920

Inspired by the solution suggested here:

import numpy_indexed as npi
ind = np.argwhere(a == a.max(0))
l = np.array(npi.group_by(ind[:,1]).split(ind[:, 0]))
def pick_one(a, index, buffer, visited):
    if index == len(a):
        return True
    for item in a[index]:
        if item not in visited:
            buffer.append(item)
            visited.add(item)
            if pick_one(a, index + 1, buffer, visited):
                return True
            buffer.pop()
            visited.remove(item)
    return False


buffer = []
pick_one(l, 0, buffer, set())
print(buffer)

example:

a = np.array([[1, 1, 0],
              [1, 0, 1],
              [0, 0, 1]])

output:

[1, 0, 2]
user_49759810
Related