Numpy выбирает определенный индекс столбца для каждой строки, используя список индексов

Я пытаюсь выбрать конкретные столбцы в строке матрицы numpy. Я после хорошего способа решить следующую проблему.

Предположим, что у меня есть следующая матрица, которую я бы назвал X

[1, 2, 3]
[4, 5, 6]
[7, 8, 9]

У меня также есть список индексов столбцов на каждую строку, которую я бы назвал Y

[1, 0, 2]

Мне нужно получить значения

[2]
[4]
[9]

Вместо списка с индексами Y я также могу создать матрицу с той же формой, что и X, где каждый столбец является bool/(int в диапазоне 0-1 значение, указывающее, является ли это требуемым столбцом.

[0, 1, 0]
[1, 0, 0]
[0, 0, 1]

Я знаю, что это можно сделать с помощью итерации по массиву и выбора значений столбца, которые мне нужны, однако мне было интересно, есть ли лучшее решение. Это будет выполняться часто на больших массивах данных, и поэтому я должен работать как можно быстрее.

Спасибо

Ответ 1

Если у вас есть логический массив, вы можете сделать прямой выбор на основе этого:

>>> a = np.array([True, True, True, False, False])
>>> b = np.array([1,2,3,4,5])
>>> b[a]
array([1, 2, 3])

Чтобы выполнить ваш первоначальный пример, вы можете сделать следующее:

>>> a = np.array([[1,2,3], [4,5,6], [7,8,9]])
>>> b = np.array([[False,True,False],[True,False,False],[False,False,True]])
>>> a[b]
array([2, 4, 9])

Вы также можете добавить arange и сделать прямой выбор, хотя в зависимости от того, как вы генерируете свой логический массив и что ваш код выглядит как YMMV.

>>> a = np.array([[1,2,3], [4,5,6], [7,8,9]])
>>> a[np.arange(len(a)), [1,0,2]]
array([2, 4, 9])

Надеюсь, что это поможет, дайте мне знать, если у вас возникнут вопросы.

Ответ 2

Вы можете сделать что-то вроде этого:

In [7]: a = np.array([[1, 2, 3],
   ...: [4, 5, 6],
   ...: [7, 8, 9]])

In [8]: lst = [1, 0, 2]

In [9]: a[np.arange(len(a)), lst]
Out[9]: array([2, 4, 9])

Подробнее об индексировании многомерных массивов: http://docs.scipy.org/doc/numpy/user/basics.indexing.html#indexing-multi-dimensional-arrays

Ответ 3

Вы можете сделать это, используя итератор. Вот так:

np.fromiter((row[index] for row, index in zip(X, Y)), dtype=int)

Время:

N = 1000
X = np.zeros(shape=(N, N))
Y = np.arange(N)

#@Aशwini चhaudhary
%timeit X[np.arange(len(X)), Y]
10000 loops, best of 3: 30.7 us per loop

#mine
%timeit np.fromiter((row[index] for row, index in zip(X, Y)), dtype=int)
1000 loops, best of 3: 1.15 ms per loop

#mine
%timeit np.diag(X.T[Y])
10 loops, best of 3: 20.8 ms per loop

Ответ 4

Как чистый numpythonic подход, вы можете получить индексы со второй оси с помощью np.take(), тогда диагональ результата будет вашим ожидаемым выходом:

np.diagonal(np.take(arr, idx, axis=1))

Демо:

>>> arr = np.array([[1, 2, 3],
... [4, 5, 6],
... [7, 8, 9]])
>>> 
>>> idx = [1, 0, 2]
>>> 
>>> np.diagonal(np.take(arr, idx, axis=1))
array([2, 4, 9])