Индексный номер nd массива вдоль последнего измерения

Есть ли простой способ индексировать многомерный массив numpy вдоль последнего измерения, используя массив индексов? Например, возьмите массив a формы (10, 10, 20). Предположим, что у меня есть массив индексов b, формы (10, 10), чтобы результат был c[i, j] = a[i, j, b[i, j]].

Я попробовал следующий пример:

a = np.ones((10, 10, 20))
b = np.tile(np.arange(10) + 10, (10, 1))
c = a[b]

Однако это не работает, потому что тогда он пытается индексироваться как a[b[i, j], b[i, j]], который не совпадает с a[i, j, b[i, j]]. И так далее. Есть ли простой способ сделать это, не прибегая к циклу?

1 ответ

Есть несколько способов сделать это. Пусть сначала сгенерированы некоторые тестовые данные:

In [1]: a = np.random.rand(10, 10, 20)
In [2]: b = np.random.randint(20, size=(10,10)) # random integers in range 0..19

Одним из способов решения вопроса было бы создать два вектора индекса, где один - вектор строки, а другой - вектор столбца 0..9, используя meshgrid:

In [3]: i1, i0 = np.meshgrid(range(10), range(10), sparse=True)
In [4]: c = a[i0, i1, b]

Это работает, потому что i0, i1 и b будут транслироваться до 10x10. Быстрый тест на правильность:

In [5]: all(c[i, j] == a[i, j, b[i, j]] for i in range(10) for j in range(10))
Out[5]: True

Другим способом было бы использовать choose и rollaxis:

# choose needs a sequence of length 20, so move last axis to front
In [22]: aa = np.rollaxis(a, -1) 
In [23]: c = np.choose(b, aa)
In [24]: all(c[i, j] == a[i, j, b[i, j]] for i in range(10) for j in range(10))
Out[24]: True

licensed under cc by-sa 3.0 with attribution.