Использование индексации целочисленных массивов для извлечения элементов на основе индексов вторичного массива

В данном сценарии цель состоит в том, чтобы получить определенные элементы из массив A с использованием индексов, указанных во втором массиве B. Вместо того, чтобы полагаться на np.take или np.choose, более простой подход — использовать NumPy. индексирование целочисленного массива:

A[np.arange(A.shape[0]),B.ravel()]

Вот как этот код достигает желаемого результата:

• np.arange(A.shape[0]) создает массив индексов от 0 до (A.shape[0] - 1), представляющий строки A.
• B.ravel() преобразует B в одномерный массив, гарантируя, что он содержит то же количество элементов, что и строки в A.
• Операция индексации A[... ] извлекает элементы из A, используя индексы строк и индексы столбцов из B.

Этот подход особенно полезен, когда B представляет собой одномерный массив или список индексов столбцов. Если пропустить операцию сглаживания, код становится еще проще:

A[np.arange(A.shape[0]),B]

Пример:

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

result = A[np.arange(A.shape[0]), B]

print(result)
# Output: [1, 2, 5]