使用索引数组在2D数组中索引列值 [英] Index column values in 2D array using array of indices

问题描述

我有以下数组:

import numpy as np 

print(A)
array([[ 0,  1,  4,  5,  8,  7],
       [ 5,  3,  4,  1,  8, 11],
       [ 2,  7,  5,  3,  4,  1],
       [ 2,  8,  8,  1, 10,  1],
       [ 2, 14,  8,  6,  5,  3]])

我需要将这些列索引对应的值A:

And I need to the values A corresponding to these column indices:

b = np.array([5, 0, 3, 4, 4])

预期输出:

array([ 7,  5,  3, 10,  5])

提前谢谢.

推荐答案

您可以使用

You can use advanced indexing. You need to define an indexing array across the first axis, so that both indexing arrays are broadcast together and each column index refers to a specific row. In this case you just want an np.arange to index on the rows:

A[np.arange(A.shape[0]), b]
# array([ 7,  5,  3, 10,  5])

这篇关于使用索引数组在2D数组中索引列值的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！