使用另一个数组作为索引从数组中选择元素

Question

说我有一个数组

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

索引数组是

B = array([[1], # want [0, 1] element of A
           [0],  # want [1, 0], element of A
           [1]])  # want [2, 1] elemtn of A

通过这个索引数组B，我想要一个3-by-1数组，其元素取自数组A，即

C = array([[2],
           [4],
           [8]])

我试过了numpy.choose，但我做不到。

Answer 1

为了答案的完整性...花哨的索引数组被广播到一个共同的形状，所以以下也可以工作，并让你最后重塑：

>>> A[np.arange(3)[:, None], B]
array([[2],
       [4],
       [8]])

Answer 2

您可以这样做：

>>> A[np.arange(len(A)), B.ravel()].reshape(B.shape)
array([[2],
       [4],
       [8]])

Answer 3

你可以这样做：

>>>np.diag(A[range(3),B]).reshape(B.shape)
array([[2],
       [4],
       [8]])

如果您想使用choose，可以执行：np.choose(B.ravel(), A.T).reshape(B.shape)。