python numpy `np.take` с 2-мерным массивом

Question

Я пытаюсь взять список элементов из массива 2D numpy с заданным списком координат, и я хочу избежать использования l oop. Я видел, что np.take работает с 1D массивом, но я не могу заставить его работать с 2D массивами.

Пример:

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]

Я хочу избежать l oop, потому что я думаю это будет медленнее (мне нужна скорость). Но если вы можете убедить меня, что все oop так же быстро, как существующее функциональное решение numpy, то я могу 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)