Python - 使用一维数组中的值作为二维数组的列索引[重复]答案

【问题标题】：Python - Use values from 1D array as column indices for 2D array [duplicate]Python - 使用一维数组中的值作为二维数组的列索引[重复]
【发布时间】：2017-04-01 22:58:39
【问题描述】：

用几句话解释它有点棘手（并谷歌它），所以：

我有这个 2D np 数组：

import numpy as np

x = np.array([[0,1,2],[3,4,5],[6,7,8],[9,10,11],[12,13,14],[15,16,17]])

还有这个一维 np 数组：

y = np.array([0,2,1,0,2,0])

我想做的是使用 y 作为（列）索引从 x 返回列值，因此它将返回如下内容：

[0, 5, 7, 9, 14, 15]

在丑陋的代码中可以这样解决：

for row,col in zip(x,y):
    print(row[col])

并且在一个不那么难看的代码中：

[row[col] for row,col in zip(x,y)]

还有其他方法可以解决这个问题吗？我想要类似的东西：

x[y]

或一个 numpy 特定的函数。

【问题讨论】：

标签： python arrays numpy

【解决方案1】：

您可以使用高级索引：

x[np.arange(6), y]
# array([ 0,  5,  7,  9, 14, 15])

【讨论】：

非常感谢，效果非常好:)

【解决方案2】：

是的，使用 numpy 会更好，因为它比使用 for 循环快几倍。我的建议是使用np.ix_() 函数对数组进行切片，然后对结果取对角线以从每一行中获取值：

np.diag(x[np.ix_(range(0, np.shape(x)[0]), y)])

【讨论】：