python numpy `np.take` 二维数组

Question

我正在尝试从具有给定坐标列表的 2D numpy 数组中获取元素列表，我想避免使用循环。 我看到np.take适用于一维数组，但我无法使其适用于二维 arrays。

例子：

a = np.array([[1,2,3], [4,5,6]])
print(a)
# [[1 2 3]
#  [4 5 6]]

np.take(a, [[1,2]])
# gives [2, 3] but I want just [6]

我想避免循环，因为我认为那会更慢（我需要速度）。 但是如果你能说服我一个循环和现有的 numpy function 解决方案一样快，那么我可以 go 解决它。

Answer 1

如果我理解正确的话，你有一个这样的坐标列表：

coords = [[y0, x0], [y1, x1], ...]

要在这些坐标处获取数组a的值，您需要：

a[[y0, y1, ...], [x0, x1, ...]]

所以a[coords]将不起作用。 一种方法是：

Y = [c[0] for c in coords]
X = [c[1] for c in coords]

要么

Y = np.transpose(coords)[0]
X = np.transpose(coords)[1]

然后

a[Y, X]

Answer 2

花式索引是否可以满足您的需求？ np.take 似乎在操作前将数组展平。

import numpy as np

a = np.arange(1, 10).reshape(3,3)

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

rows = [ 1,1,2,0]
cols = [ 0,1,1,2]

# Use the indices to access items in a
a[rows, cols]
# array([4, 5, 8, 3])

a[1,0], a[1,1], a[2,1], a[0,2]
# (4, 5, 8, 3)