使用numpy数组作为另一个数组的第二个dim的索引？ [英] using an numpy array as indices of the 2nd dim of another array?

问题描述

例如，我有两个numpy数组，

For example, I have two numpy arrays,

A = np.array(
  [[0,1], 
   [2,3], 
   [4,5]])
B = np.array(
  [[1],
   [0],
   [1]], dtype='int')

我要提取 A 的每一行中的一个元素，该元素由 B 索引，所以我想要以下结果：

and I want to extract one element from each row of A, and that element is indexed by B, so I want the following results:

C = np.array(
  [[1],
   [2],
   [5]])

我试过 A [：，B.ravel（ ）] ，但它会广播 B ，而不是我想要的。还看了 np.take ，似乎不是解决我问题的正确方法。

I tried A[:, B.ravel()], but it'll broadcast B, not what I want. Also looked into np.take, seems not the right solution to my problem.

但是，我可以用 np.choose 通过转置 A ，

However, I could use np.choose by transposing A,

np.choose(B.ravel(), A.T)

但任何其他更好的解决方案？

but any other better solution?

推荐答案

你可以使用 NumPy的纯整数数组索引 -

You can use NumPy's purely integer array indexing -

A[np.arange(A.shape[0]),B.ravel()]

样品运行 -

In [57]: A
Out[57]: 
array([[0, 1],
       [2, 3],
       [4, 5]])

In [58]: B
Out[58]: 
array([[1],
       [0],
       [1]])

In [59]: A[np.arange(A.shape[0]),B.ravel()]
Out[59]: array([1, 2, 5])

请注意，如果 B 是 1D 数组或此类列索引的列表，您可以使用<$ c $跳过展平操作c> .ravel（）。

Please note that if B is a 1D array or a list of such column indices, you could simply skip the flattening operation with .ravel().

示例运行 -

In [186]: A
Out[186]: 
array([[0, 1],
       [2, 3],
       [4, 5]])

In [187]: B
Out[187]: [1, 0, 1]

In [188]: A[np.arange(A.shape[0]),B]
Out[188]: array([1, 2, 5])

这篇关于使用numpy数组作为另一个数组的第二个dim的索引？的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！