在numpy中通过另一个数组索引一个数组 [英] Indexing one array by another in numpy

问题描述

假设我有一个带有一些任意值的矩阵 A:

Suppose I have a matrix A with some arbitrary values:

array([[ 2, 4, 5, 3],
       [ 1, 6, 8, 9],
       [ 8, 7, 0, 2]])

以及包含 A 中元素索引的矩阵 B:

And a matrix B which contains indices of elements in A:

array([[0, 0, 1, 2],
       [0, 3, 2, 1],
       [3, 2, 1, 0]])

如何从B指向的A中选择值，即:

How do I select values from A pointed by B, i.e.:

A[B] = [[2, 2, 4, 5],
        [1, 9, 8, 6],
        [2, 0, 7, 8]]

推荐答案

np.take_along_axis 是自 numpy 1.15 以来实现的此用例的内置函数.请参阅下方@hpaulj 的答案以了解如何使用它.

np.take_along_axis is a builtin function for this use case implemented since numpy 1.15. See @hpaulj 's answer below for how to use it.

您可以使用 NumPy 的高级索引 -

You can use NumPy's advanced indexing -

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

也可以使用线性索引 -

m,n = A.shape
out = np.take(A,B + n*np.arange(m)[:,None])

样品运行 -

In [40]: A
Out[40]: 
array([[2, 4, 5, 3],
       [1, 6, 8, 9],
       [8, 7, 0, 2]])

In [41]: B
Out[41]: 
array([[0, 0, 1, 2],
       [0, 3, 2, 1],
       [3, 2, 1, 0]])

In [42]: A[np.arange(A.shape[0])[:,None],B]
Out[42]: 
array([[2, 2, 4, 5],
       [1, 9, 8, 6],
       [2, 0, 7, 8]])

In [43]: m,n = A.shape

In [44]: np.take(A,B + n*np.arange(m)[:,None])
Out[44]: 
array([[2, 2, 4, 5],
       [1, 9, 8, 6],
       [2, 0, 7, 8]])

这篇关于在numpy中通过另一个数组索引一个数组的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！