问题描述
例如,我有两个 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.
但是,我可以通过转置 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
数组或此类列索引的列表,您可以简单地使用 .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 数组作为另一个数组的第二个维度的索引?的文章就介绍到这了,希望我们推荐的答案对大家有所帮助,也希望大家多多支持!