如何索引numpy数组?
问题描述:
一个简单的例子:我有一个numpy数组([[5,4,2],[1,8,8],[4,5,6]])和另一个数组([2,2,1]) 。第二个数组包含第一个数组中每行所关注的索引。我想我的最终数组只是一个矢量数组([2,8,5])。我正在寻找一些沿着first_array [second_array]的线路,但它证明难以捉摸。谢谢!如何索引numpy数组?
答
你可以尝试像以下:
>>> arr
array([[5, 4, 2],
[1, 8, 8],
[4, 5, 6]])
>>> idx
array([2, 2, 1])
>>> arr[np.arange(arr.shape[0]), idx]
array([2, 8, 5])