python - 根据另一个数组中指定的列从二维 numpy 数组中获取每行的值

Question

给定一个 2D numpy 数组A和一个 1D numpy 数组col，它们具有相同的行数和列的值A，我怎样才能得到一个新的 1D 数组，它从 中的每一行中选择相应的列A？

例子：

np.random.seed(0)
A = np.random.randint(10, size=(5, 3))
A
# array([[5, 0, 3],
#        [3, 7, 9],
#        [3, 5, 2],
#        [4, 7, 6],
#        [8, 8, 1]])
col = np.random.randint(A.shape[1], size=A.shape[0])
col
# array([2, 2, 0, 1, 1])

基于col，我想从第 0 行和第 1 行获取元素 2，从第 2 行获取元素 0，从第 3 行和第 4 行获取元素 1: [3, 9, 3, 7, 8]。

正常的索引不起作用：

A[col]
# array([[3, 5, 2],
#        [3, 5, 2],
#        [5, 0, 3],
#        [3, 7, 9],
#        [3, 7, 9]])
A[:, col]
# array([[3, 3, 5, 0, 0],
#        [9, 9, 3, 7, 7],
#        [2, 2, 3, 5, 5],
#        [6, 6, 4, 7, 7],
#        [1, 1, 8, 8, 8]])

标签： pythonnumpy

Answer 1

这可行，但对于大型数组来说非常慢：

A[:, col].diagonal()
# array([3, 9, 3, 7, 8])