二维numpy数组与另一个数组的切片

fivyi3re 于 2022-11-23 发布在其他

关注(0)|答案(1)|浏览(185)

我有一个很大的二维数组，我需要得到轴1中最大值的索引和最大值本身。我可以如下检索这两个值：

import numpy as np
a = np.arange(27).reshape(9, 3)
idx = np.argmax(a, axis=1)
max_val = np.max(a, axis=1)

但是，由于我已经找到了最大值的索引，所以感觉我应该能够使用idx构造最大值的数组，而不必再次查找该值。
我知道我可以使用np.choose(idx, a.T)，但这涉及到转置矩阵，这将比仅仅使用max要昂贵得多。我可以做一些类似np.array([a[i][idx[i]] for i in range(len(a))])的事情，但这涉及到创建一个列表，这似乎也比仅仅调用np.max要昂贵得多。
有没有办法在numpy中将a与idx切片而不重新构造数组？

numpy

来源：https://stackoverflow.com/questions/74495151/slice-of-2d-numpy-array-with-another-array

1条答案

按热度按时间

h43kikqp1#

您的a和argmax：

In [602]: a
Out[602]: 
array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11],
       [12, 13, 14],
       [15, 16, 17],
       [18, 19, 20],
       [21, 22, 23],
       [24, 25, 26]])

In [603]: idx
Out[603]: array([2, 2, 2, 2, 2, 2, 2, 2, 2], dtype=int64)

使用索引数组的常见方法：

In [606]: a[np.arange(a.shape[0]),idx]
Out[606]: array([ 2,  5,  8, 11, 14, 17, 20, 23, 26])

一个更新的工具，可能更容易使用（如果不熟悉第一个）：

In [607]: np.take_along_axis(a,idx[:,None],1)
Out[607]: 
array([[ 2],
       [ 5],
       [ 8],
       [11],
       [14],
       [17],
       [20],
       [23],
       [26]])

赞(0）回复(0）举报 2022-11-23

我来回答

二维numpy数组与另一个数组的切片

1条答案

相关问题

热门标签

最新问答