有没有一种方法可以找到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

下面是一个稍微复杂一点的例子：

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

在这种情况下，没有一个唯一的最大值列。我会满足于以下任何一种答案：

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

甚至更复杂的例子是：

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

在这种情况下，我会对这些答案的任意一个表示满意：

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

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

- ToddP

---

这个最大值总是1吗？ - amzon-ex

不，最大值可以是任何值，在每个列中也可以有不同的最大值。但现在，我会满足于上述情况的解决方案。 - ToddP

这种最大位置的结构是否存在保证（指每列中都有行索引承载其最大值，可能与其他列不同）？ - Ehsan

是的，就像上面的情况一样，有这个保证。 - ToddP

@ToddP 如果你的矩阵不是太大（或每个数组中的最大值没有重复太多次），下面建议的答案可能会有所帮助。它也适用于各列中的不同最大值。 - Ehsan

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)

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

或者，查看此帖子以获取图论解决方案。

- Paul Panzer

---

非常棒的答案。已点赞。随时愿意采纳你的答案 :) 我最初发布问题是为了看看是否有更好的答案，但似乎这个答案甚至比那个更好。 - Ehsan

Answer 2

受这里提出的解决方案启发：

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)

例子：

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

输出：

[1, 0, 2]

- Ehsan

---

这很好，但请看我的原始帖子的结尾：我正在寻找一种不需要循环和逻辑检查的解决方案。 - ToddP

我理解你的观点。虽然我不知道@Paul提供的scipy包后端，但我敢打赌它比我的答案更快。祝好运。 - Ehsan

@ToddP 另外，如果你更关心代码的性能而不是可读性，我猜使用Paul在我的帖子https://stackoverflow.com/questions/62571292/find-a-list-of-unique-representatives-elements-from-a-list-of-arrays中的答案会更快。随意查看。`l`在我的答案中将是要提供给该帖子答案的数组。如果您需要帮助合并两个答案，请告诉我。 - Ehsan