有没有办法找到 2D NumPy 数组中最大列值的唯一行索引？

Question

对于 2D NumPy 数组中的每一列，该列的最大值可以出现多次。我想找到每列最大值的行索引，而不重复行索引。

np.argmax这是一个演示为什么不起作用的示例：

import numpy as np

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

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

print(ind)

输出：

[0 0 2]

我想要结果：[1, 0, 2]对于这个例子。

那是：

• 第二列的行索引必须为 0
• 这意味着第一列的行索引必须为 1
• 这又意味着第三列的行索引必须为 2

一个稍微复杂一点的例子是这个数组：

[0 0 2]

在这种情况下，不存在具有唯一最大值的列。我会对以下任一答案感到满意：

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

一个更复杂的例子是：

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

在这种情况下，我会对以下任何答案感到满意：

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

我可以用循环和逻辑条件解决这些问题，但我想知道是否有办法使用 numpy 函数解决问题？

Answer 1

可能有点矫枉过正，但你可以使用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])

请注意，您始终可以使用类似的方法减少到 0,1 情况

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

这可以大大加快分配速度。

或者，请参阅这篇文章以获取图论解决方案。