numpy 使用在另一个ndarray中定义的索引切片ndarray

ctehm74n 于 2023-08-05 发布在其他

关注(0)|答案(2)|浏览(98)

假设我有以下内容：

a = np.array([[1,2,3],[4,5,6]])
idx = np.array([[0,2],[1,2]])

字符串
有没有一种方法可以对数组进行“numpy-slice”，这样我就可以使用idx中的第i行作为a中第i行的索引，即a[i,idx[i]]？
这可以很容易地通过一个循环来完成：

np.vstack([a[i, idx[i]] for i in range(len(a))])

型
但我想一定有一种 numpy 方式。
最终结果应该是：

[[1,3],
 [5,6]]

型
做a[idx]不起作用。如果我做a[:,idx]，那么我会得到太多的结果，即给出

array([[[1, 3],
        [2, 3]],
       [[4, 6],
        [5, 6]]])

型
很接近但不正确

2条答案

使用np.take_along_axis沿着1轴选择所需的索引：

>>> np.take_along_axis(a, idx, 1)
array([[1, 3],
       [5, 6]])

字符串

你可以使用NumPy的高级索引来获得你想要的结果。

import numpy as np

a = np.array([[1,2,3],[4,5,6]])
idx = np.array([[0,2],[1,2]])
res = a[np.arange(a.shape[0])[:,None], idx]
print(res)

字符串
输出量：

[[1 3]
 [5 6]]

型