Python - 使用1D数组中的值作为2D数组的列索引

Question

用很少的单词（和谷歌）解释它很难，所以：

我有这个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]])

和这个1D 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特定功能。

Answer 1

您可以使用高级索引：

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

Answer 2

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

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