首页 文章

使用numpy数组作为另一个数组的第二个dim的索引？ [重复]

提问于
浏览 次
3

这个问题在这里已有答案：

例如，我有两个numpy数组，

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

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

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

我试过 A[:, B.ravel()] ，但它会播出 B ，而不是我想要的 . 也看了 np.take ，似乎不是我的问题的正确解决方案 .

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

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

但还有其他更好的解决方案

python numpy indexing

2 回答

  • 3

    你可以使用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 数组或此类列索引的列表，则可以使用 .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])
    回复于
  • 0
    C = np.array([A[i][j] for i,j in enumerate(B)])
    回复于

相关问题