python - 从另一个数组中选择具有列索引的每一行中的元素

Question

正如标题所说，例如，给定一个 2d numpy 数组

a = np.array([[2,3,4,5],[12,4,5,7],[14,2,5,6],[12,3,4,6]])
idx = np.array([[0,2],[2,3],[1,3],[1,3]])

我想从第一行中选择第一个和第三个元素，2、4等。所以最终的答案应该是

ans = [[2,4],
       [5,7],
       [2,6],
       [3,6]]

我已经尝试过 np.choose 和 np.take 但我相信 np.take 会使数组变平，并且 np.choose 看起来不像我预期的那样。

任何想法将不胜感激！非常感谢！

Answer 1

使用List comprehension，您可以遍历 2 numpy 数组并创建一个具有所需输出的新列表。稍后，您将每 2 个元素添加到一个 numpy 数组中，这可以使用以下reshape函数来实现：

aa = [a[index][x] for index, value in enumerate(idx) for x in value]
# aa = [2, 4, 5, 7, 2, 6, 3, 6]
aa = np.reshape(aa, (-1, 2))
print(aa)
# output 
[[2 4]
 [5 7]
 [2 6]
 [3 6]]

Answer 2

您可以定义一些复杂的函数将您idx转换为切片，但我认为通过列表理解和枚举访问数组会更容易，然后将列表转换为数组。像这样：

a = np.array([[2,3,4,5],[12,4,5,7],[14,2,5,6],[12,3,4,6]])
idx = np.array([[0,2],[2,3],[1,3],[1,3]])

np.array([a[row,elems] for row,elems in enumerate(idx)])

结果：

array([[2, 4],
       [5, 7],
       [2, 6],
       [3, 6]])

Answer 3

您可以在 numpy 中使用高级索引：

a[np.indices(idx.shape)[0], idx]

np.indices(idx.shape)[0]为您的列索引创建相应的行索引idx，它们形成高级索引。

输出：

[[2 4]
 [5 7]
 [2 6]
 [3 6]]