使用二维数组作为指示符从二维numpy数组中提取元素

Question

我有一个二维数组：

>>> input = np.array([[1,2],[3,4]])
>>> input
array([[1, 2],
       [3, 4]])

我还有一个数组，它的列数与输入数组的列数相同，每行都包含要从输入数组的匹配列中提取的索引。你知道吗

>>> indices = np.array([[0],[1]])
>>> indices
array([[0],
       [1]])

在本例中，我希望获得以下数组作为输出：

array([[1],
       [4]])

我有办法做到吗？你知道吗

Answer 1

您可以先对索引进行切片：

in_array=np.arange(20).reshape(4,5)
ind=np.array([[0,2],[1,3]])
in_array[ind[:,0],ind[:,1]]
array([2, 8])

也不要使用诸如input之类的名称，因为它们是python的内置函数，您正在覆盖它们。你知道吗

Answer 2

你可以试试这个：

#!python
# -*- coding: utf-8 -*-#
#
# Imports
import numpy as np

arr = np.array([[1,2],[3,4]])
idx = np.array([[0],[1]])

ans = arr[idx,idx]
print(ans)

这回答了OP的具体问题，但是，如果我们想让它更一般化，请遵循@jpp的思想，并这样做：

# Imports
import numpy as np

arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
idx = np.array([0, 1, 1])

col_vec1 = np.arange(arr.shape[0])[:, None]
col_vec2 = idx[:, None]

ans = arr[col_vec1, col_vec2]

print(arr[0, 0])  # 1
print(arr[1, 1])  # 5
print(arr[2, 1])  # 8

print(ans)

Answer 3

单程通过NumPy数组索引。你知道吗

A = np.array([[1,2],[3,4]])
idx = np.array([[0],[1]])

res = A[np.arange(A.shape[0])[:, None], idx]

print(res)

[[1]
 [4]]

注意，我们需要确保行和列索引器具有相同的形状，在本例中(n, 1)，其中n是行数。你知道吗