Numpy Выбор элементов по заданным индексам массивов строк и столбцов

Question

У меня есть индексы строк в виде массива 1d numpy и список массивов numpy (список такой же длины, что и размер массива индексов строк. Я хочу извлечь значения, соответствующие этим индексам. Как я могу сделать это?

Это пример того, что я хочу в качестве вывода с учетом ввода

A = np.array([[2, 1, 1, 0, 0],
              [3, 0, 2, 1, 1],
              [0, 0, 2, 1, 0],
              [0, 3, 3, 3, 0],
              [0, 1, 2, 1, 0],
              [0, 1, 3, 1, 0],
              [2, 1, 3, 0, 1],
              [2, 0, 2, 0, 2],
              [3, 0, 3, 1, 2]])

row_ind = np.array([0,2,4])
col_ind = [np.array([0, 1, 2]), np.array([2, 3]), np.array([1, 2, 3])]

Теперь я хочу, чтобы мои выходные данные представляли собой список numpy массивов или список списков как

[np.array([2, 1, 1]), np.array([2, 1]), np.array([1, 2, 1])]

Больше всего меня беспокоит эффективность. Массив A имеет размер 20K x 10K.

Answer 1

Как прокомментировал @hpaulj, вероятно, вы не сможете избежать зацикливания - например,

import numpy as np

A = np.array([[2, 1, 1, 0, 0],
              [3, 0, 2, 1, 1],
              [0, 0, 2, 1, 0],
              [0, 3, 3, 3, 0],
              [0, 1, 2, 1, 0],
              [0, 1, 3, 1, 0],
              [2, 1, 3, 0, 1],
              [2, 0, 2, 0, 2],
              [3, 0, 3, 1, 2]])


row_ind = np.array([0,2,4])
col_ind = [np.array([0, 1, 2]), np.array([2, 3]), np.array([1, 2, 3])]

# make sure the following code is safe...
assert row_ind.shape[0] == len(col_ind)

# 1) select row (A[r, :]), then select elements (cols) [col_ind[i]]:
output = [A[r, :][col_ind[i]] for i, r in enumerate(row_ind)]

# output
# [array([2, 1, 1]), array([2, 1]), array([1, 2, 1])] 

Другим способом сделать это может быть использование np.ix_ (по-прежнему требуется зацикливание). Используйте с осторожностью, хотя для очень больших массивов; np.ix_ использует расширенное индексирование - в отличие от базового среза c, оно создает copy данных вместо view - см. документы .