在二维数组中查找 n 个 argmins

Posted 2023-03-17

【中文标题】在二维数组中查找 n 个 argmins

【英文标题】：Find n argmins in 2d array

【发布时间】：2021-12-20 18:53:08

【问题描述】：

我需要在 python 的二维数组中找到 n 个 argmins，我该怎么做？

示例：

a=np.array([[1, 5, 9], [2, 3, 10], [4, 11, 12]])
argmins = n_argmins(a, 5) # finding 5 argmins
print(argmins)

输出：

[[0 1]
 [0 2]
 [1 0]
 [1 1]
 [2 0]]

【问题讨论】：

结果不应该是这样吗？ [[0, 0] [1, 0] [1, 1] [2, 0] [0, 1]]

您可以从argsort 或a.ravel() 开始

【参考方案1】：

您可以尝试以下方法：

import numpy as np

a=np.array([[1, 5, 9], [2, 3, 10], [4, 11, 12]])
np.c_[np.unravel_index(a.ravel().argsort()[:5], a.shape)]

它给出：

array([[0, 0],
       [1, 0],
       [1, 1],
       [2, 0],
       [0, 1]])

【参考方案2】：

将数据转换为列表，但请记住原始索引。 在列表中搜索 n min args。

代码：

def n_argmins(a, n)
  for i,inner in enumerate(a):
     for j,x in enumerate(inner):
         indexes[x] = [i,j]
         a_new.append(x)
  res = []
  for _ in range(n):
     res.append(indexes[min(a_new)])
     a_new.remove(min(a_new))
  return res

假设所有列表都已排序（如示例所示），我们可以以更好的方式解决它。该解决方案基于每个内部列表中的第一个值已经是该列表的最小参数的想法。所以我们只需要搜索它们。代码：

n = 5
next_possible_mins = [0]*len(a)
indexes = [0]*len(a)
result = []
for i in range(len(a)):
    next_possible_mins[i] = a[i][0]
    indexes[i] = [0, i]

for i in range(n):
    index_min = next_possible_mins.index(min(next_possible_mins))
    result.append(indexes[index_min])
    indexes[index_min] = [indexes[index_min][0]+1, indexes[index_min][1]]
    next_possible_mins[index_min] = a[indexes[index_min][1]][indexes[index_min][0]]

print(result)