Question

给定一个3维numpy数组，如何找到前n个最小值的索引？最小值的索引可以是：

i,j,k = np.where(my_array == my_array.min())

Answer 1

这是通用n-dims和通用N个最小数字的一种方法 -

def smallestN_indices(a, N):
    idx = a.ravel().argsort()[:N]
    return np.stack(np.unravel_index(idx, a.shape)).T

2D输出数组的每一行都将保存与最小数组编号之一对应的索引元组。

我们也可以使用argpartition，但这可能无法维持订单。因此，我们需要在argsort那里进行更多的额外工作 -

def smallestN_indices_argparitition(a, N, maintain_order=False):
    idx = np.argpartition(a.ravel(),N)[:N]
    if maintain_order:
        idx = idx[a.ravel()[idx].argsort()]
    return np.stack(np.unravel_index(idx, a.shape)).T

示例运行 -

In [141]: np.random.seed(1234)
     ...: a = np.random.randint(111,999,(2,5,4,3))
     ...: 

In [142]: smallestN_indices(a, N=3)
Out[142]: 
array([[0, 3, 2, 0],
       [1, 2, 3, 0],
       [1, 2, 2, 1]])

In [143]: smallestN_indices_argparitition(a, N=3)
Out[143]: 
array([[1, 2, 3, 0],
       [0, 3, 2, 0],
       [1, 2, 2, 1]])

In [144]: smallestN_indices_argparitition(a, N=3, maintain_order=True)
Out[144]: 
array([[0, 3, 2, 0],
       [1, 2, 3, 0],
       [1, 2, 2, 1]])

运行时测试 -

In [145]: a = np.random.randint(111,999,(20,50,40,30))

In [146]: %timeit smallestN_indices(a, N=3)
     ...: %timeit smallestN_indices_argparitition(a, N=3)
     ...: %timeit smallestN_indices_argparitition(a, N=3, maintain_order=True)
     ...: 
10 loops, best of 3: 97.6 ms per loop
100 loops, best of 3: 8.32 ms per loop
100 loops, best of 3: 8.34 ms per loop

找到3维numpy数组中n个最小值的索引

1 个答案: