Как создать массив 3D (shape = m, n, o) с единицами вдоль определенной (третьей) оси, индексы задаются массивом 2D (shape = m, n)?

Question

Допустим, у меня есть двумерный массив с формой (2, 2), содержащий индексы

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

Что я хотел бы сделать, это создать трехмерный массив с формой (2, 2, 4), который имеет значения 1 вдоль третьей оси, и их положение задается как x, поэтому:

y = np.zeros(shape=(2,2,4))
myfunc(array=y, indices=x, axis=2)

array([[[0, 0, 1, 0],
        [1, 0, 0, 0]],
       [[0, 0, 0, 1],
        [0, 1, 0, 0]]])

Пока что я не нашел никакого метода индексации для этого. Цикл for может быть в состоянии сделать это, но я уверен, что есть более быстрый, векторизованный метод.

Answer 1

То, что вы ищете, называется расширенное индексирование .Чтобы правильно индексировать целочисленные массивы, необходимо иметь набор массивов, которые передают правильную форму.Поскольку x уже выровнено с двумя измерениями, вам нужно только создать двумерные массивы с индексами вдоль каждой оси.np.ogrid помогает в этом, поскольку он создает массивы минимального диапазона, которые передают правильную форму:

a, b = np.ogrid[:2, :2]
y[a, b, x] = 1

Результаты ogrid эквивалентны

a = np.arange(2).reshape(-1, 1)
b = np.arange(2).reshape(1, -1)

Или

a = np.arange(2)[:, None]
b = np.arange(2)[None, :]

Вы также можете написать однострочник:

y[(*tuple(slice(None, n) for n in x.shape), x)] = 1

Answer 2

С for петлями:

y = np.zeros(shape=(2,2,4))
for i in range(2):
    for j in range(2):
        y[i,j,x[i,j]] = 1

С расширенным индексированием:

y = np.zeros(shape=(2,2,4))
i, j = np.meshgrid(np.arange(2), np.arange(2))
y[i,j,x[i,j]] = 1

Выход (в обоих случаях):

array([[[0., 0., 1., 0.],
        [1., 0., 0., 0.]],
       [[0., 0., 0., 1.],
        [0., 1., 0., 0.]]])

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