Сложная индексация numpy

#python #arrays #numpy

#python #массивы #numpy

Вопрос:

У меня есть значения в 2D-массиве с именем a . У меня есть второй массив, C , в котором у меня есть индексы, которые ссылаются на ячейки в a . Следовательно, размеры этого массива являются (M, N, 2) .

Я хочу создать вывод измерений, (M, N) который состоит из значений a из индексов, хранящихся в C .

Это pythonic способ сделать это:

 a = np.arange(20).reshape(-1, 5)
print(a)
# [[ 0  1  2  3  4]
#  [ 5  6  7  8  9]
#  [10 11 12 13 14]
#  [15 16 17 18 19]]

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


output = np.array([[a[cell[0], cell[1]] for cell in row] for row in C])
print(output)

# [[10 16]
#  [ 0  9]
#  [12  6]
#  [ 0  6]]
  

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

Есть идеи?

Приветствия

Комментарии:

1. a[C[...,0], C[...,1]] ?

2. @Divakar Я никогда не перестаю удивляться хитрости индексации Numpy…

3. @Divakar Пожалуйста, опубликуйте это в качестве ответа.

Ответ №1:

@Divakar: Это действительно решение. Я так поражен… Событие не знало, что это был действительный код. Большое спасибо.

 >>> output = a[C[..., 0], C[..., 1]]
>>> output
array([[10, 16],
       [ 0,  9],
       [12,  6],
       [ 0,  6]])