我有一个numpy数组列表,每个数组都是相同的形状。让我们说:
a = [np.array([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]]),
np.array([[11, 12, 13],
[14, 15, 16],
[17, 18, 19]]),
np.array([[99, 98, 97],
[96, 95, 94],
[93, 92, 91]])]
我有另一个相同形状的数组,给出了我想从中获取元素的列表索引:
b = np.array([[0, 0, 1],
[2, 1, 0],
[2, 1, 2]])
我想得到的是以下内容:
np.array([[1, 2, 13],
[96, 15, 6],
[93, 18, 91]])
有一个简单的解决方案工作得很好:
np.choose(b, a)
但这最多只限于32个阵列。但就我而言,我必须组合更多的阵列(超过100个)。所以我需要另一种方法来做到这一点。
我想,它必须是关于进步索引或者np.take
方法的东西。所以可能,第一步是a = np.array(a)
,然后像a[np.arange(a.shape[0]), b]
。但是我没有让它发挥作用。
有人可以帮忙吗? :)
In [129]: a = [np.array([[1, 2, 3],
...: [4, 5, 6],
...: [7, 8, 9]]),
...: np.array([[11, 12, 13],
...: [14, 15, 16],
...: [17, 18, 19]]),
...: np.array([[99, 98, 97],
...: [96, 95, 94],
...: [93, 92, 91]])]
In [130]: b = np.array([[0, 0, 1],
...: [2, 1, 0],
...: [2, 1, 2]])
In [131]:
In [131]: A = np.array(a)
In [132]: A.shape
Out[132]: (3, 3, 3)
您想使用b
来索引第一个维度。对于其他维度,您需要使用b
广播的索引,即列向量和行向量:
In [133]: A[b, np.arange(3)[:,None], np.arange(3)]
Out[133]:
array([[ 1, 2, 13],
[96, 15, 6],
[93, 18, 91]])
存在用于创建这些阵列的各种便利功能,例如,
In [134]: np.ix_(range(3),range(3))
Out[134]:
(array([[0],
[1],
[2]]), array([[0, 1, 2]]))
和ogrid
在另一个答案中提到。
这是一个相对较新的功能,也可以完成这项工作:
In [138]: np.take_along_axis(A, b[None,:,:], axis=0)
Out[138]:
array([[[ 1, 2, 13],
[96, 15, 6],
[93, 18, 91]]])
在我调整b
之前,我不得不考虑一下。