array([[ 2, 4, 5, 3],
[ 1, 6, 8, 9],
[ 8, 7, 0, 2]])
还有一个矩阵B，其中包含A中元素的索引：
array([[0, 0, 1, 2],
[0, 3, 2, 1],
[3, 2, 1, 0]])
如何从B指向的A中选择值，即：
A[B] = [[2, 2, 4, 5],
[1, 9, 8, 6],
[2, 0, 7, 8]]
#1 楼
您可以使用NumPy's advanced indexing-A[np.arange(A.shape[0])[:,None],B]
也可以使用
linear indexing-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]])
#2 楼
较新的版本添加了可以执行此功能的take_along_axis函数：In [203]: A = np.array([[ 2, 4, 5, 3],
...: [ 1, 6, 8, 9],
...: [ 8, 7, 0, 2]])
In [204]: B = np.array([[0, 0, 1, 2],
...: [0, 3, 2, 1],
...: [3, 2, 1, 0]])
In [205]: np.take_along_axis(A,B,1)
Out[205]:
array([[2, 2, 4, 5],
[1, 9, 8, 6],
[2, 0, 7, 8]])
还有一个
put_along_axis。#3 楼
我知道这是一个老问题，但是使用索引的另一种方法是：A[np.indices(B.shape)[0], B]
输出：
[[2 2 4 5]
[1 9 8 6]
[2 0 7 8]]
#4 楼
以下是使用for循环的解决方案：outlist = []
for i in range(len(B)):
lst = []
for j in range(len(B[i])):
lst.append(A[i][B[i][j]])
outlist.append(lst)
outarray = np.asarray(outlist)
print(outarray)
上面也可以用更简洁的列表理解形式编写：
outlist = [ [A[i][B[i][j]] for j in range(len(B[i]))]
for i in range(len(B)) ]
outarray = np.asarray(outlist)
print(outarray)
输出：
[[2 2 4 5]
[1 9 8 6]
[2 0 7 8]]
评论
问题的追寻者正在寻找一种使用简单切片而不是循环的解决方案。切片比循环快得多。
–loganjones16
19 Mar 7 '19 at 5:05
评论
我不知道这些是如何工作的...您能补充一些解释吗？
–Jason S
17-10-31在21:38
@JasonS基本上，沿第一个轴使用范围数组为B中的每个元素沿该轴选择一个。提供的链接应提供所需的所有详细信息。
– Divakar
17-10-31在21:41
嗯...很好，但是我看了高级索引链接，对我来说这没有意义。我需要将您的示例扩展到仅应用最后一个索引的N维数组，并且我不确定该怎么做。
–Jason S
17-10-31在21:44