Numpy矩阵：根据列表选择列答案

【问题标题】：Numpy matrix: select column based on listNumpy矩阵：根据列表选择列
【发布时间】：2020-04-15 12:48:13
【问题描述】：

我有一个 numpy 数组和一个列表，该列表定义了我要选择的行。执行此操作的最佳方法是什么？

import numpy as np

a = np.array([[1,2,3],
              [4,5,6],
              [7,8,9]])

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

想要的结果

np.array([[2],
         [4],
         [9]])

我试过 np.take() 但这不起作用。

亲切的问候

编辑：由于需要在大型数组上重复执行此操作，我正在寻找一种矢量化方法（无循环）

【问题讨论】：

实现这个的数学逻辑是什么？
对于a 中的每一行，您选择b 定义的列中的元素
因此您想要a 的第一行的第一个元素、第二行的第0 个元素和最后一行的第二个元素（定义在b）
这能回答你的问题吗？ NumPy selecting specific column index per row by using a list of indexes

标签： python arrays numpy

【解决方案1】：

如果您从b 中删除无关维度

b = np.sqeeze(b)

您可以使用以下内容：

a[np.arange(len(b)), b]

【讨论】：

看看新的np.take_along_axis

【解决方案2】：

这不是很pythonic，但这应该可以解决您的问题：

res = np.zeros(len(b))
for i, row in enumerate(a):
    res[i] = row[b[i]]

print(res)

一行相同：

a[[i[0] for i in b],[i for i in range(len(b))]]

【讨论】：

谢谢，这确实是我们想要的行动！有没有办法不用循环？

【解决方案3】：

最近的版本添加了一个take_along_axis，它可以满足您的需求：

In [96]: a = np.array([[1,2,3], 
    ...:               [4,5,6], 
    ...:               [7,8,9]]) 
    ...:  
    ...: b = np.array([[1], 
    ...:               [0], 
    ...:               [2]])                                                                           
In [97]: np.take_along_axis(a,b,axis=1)                                                                
Out[97]: 
array([[2],
       [4],
       [9]])

它的工作原理很像@Nils answer、a[np.arange(3), np.squeeze(b)]，但处理尺寸更好。