python - python numpy`np.take`与二维数组

Question

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

例子：

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 函数解决方案一样快，那么我可以选择它。

标签： pythonnumpy

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]