花哨的索引一个 numpy 矩阵：每行一个元素

Question

我有一个形状为 (m, n) 的二维 numpy 数组、矩阵。我的实际用例有 m ~ 1e5 和 n ~ 100，但为了有一个简单的最小示例：

matrix = np.arange(5*3).reshape((5, 3))

我有一个整数索引数组 idx，形状为 (m, )，每个条目都在 [0, n) 之间。此数组指定应从matrix 的每一行中选择哪一列。

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

所以，我试图从第 0 行选择第 2 列，从第 1 行选择第 0 列，从第 2 行选择第 2 列，从第 1 行选择第 1 列，从第 4 行选择第 1 列。因此最终答案应该是：

correct_result = np.array((2, 3, 8, 10, 13))

我已经尝试了以下，这是直观的，但不正确：

incorrect_result = matrix[:, idx]

上面的语法所做的是将 idx 逐行应用为花哨的索引数组，从而产生另一个形状为 (m, n) 的矩阵，这不是我想要的。

这种类型的花式索引的正确语法是什么？

Answer 1

correct_result = matrix[np.arange(m), idx]

高级索引表达式matrix[I, J]给出的输出使得output[n] == matrix[I[n], J[n]].

如果我们想要output[n] == matrix[n, idx[n]]，那么我们需要I[n] == n和J[n] == idx[n]，所以我们需要I成为np.arange(m)和J成为idx。