Эффективная многомерная индексация в numpy

Question

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

Итак, давайте начнем с примера:

# The array I would like to sort
A = np.ones([3, 10, 2])

# The index array
i = np.ones([10, 2], dtype=int)

# Bring some life into sample data
samples = np.arange(10, dtype=int)
A[2, :, 0] = A[2, :, 0] * samples
np.random.shuffle(samples)
A[2, :, 1] = A[2, :, 1] * samples  
i[:, 0] = i[:, 0] * samples
np.random.shuffle(samples)
i[:, 1] = i[:, 1] * samples

Итак, мой массив A содержит 2 ломтика из 10 наборов по 3 значениям. То, что я хочу сделать, это отсортировать каждый фрагмент по отдельности, сохраняя каждый набор вместе. Имея индексный массив i, мое решение:

A = A[:, i]
shape = A.shape
A = A.reshape([shape[0], shape[1], shape[2] + shape[3]])
A = A[:, :, [0, 3]]

, где я впервые использую i для индексации A. Это создает новое измерение, где каждый столбец i применяется к A и заканчивается массивом формы (4, 10, 2, 2). Поскольку мне нужно только два из четырех результатов, я изменяю массив и удаляю ненужную мне информацию.

Этот подход работает нормально, но мне интересно, есть ли более эффективное или элегантное решение для этого.

Юльз

Answer 1

Вы можете использовать advanced indexing здесь:

A[:,i, np.arange(A.shape[2])]

---

Сравнение с текущим подходом -

out = A[:, i]
shape = out .shape
out  = out .reshape([shape[0], shape[1], shape[2] + shape[3]])
out  = out [:, :, [0, 3]]

np.allclose(A[:,i, np.arange(A.shape[2])], out)
#True