Блог

Имея массив A с формой (2,6, 60) , возможно ли его индексировать на основе двоичного массива B формы (6,) ?

6 и 60 довольно произвольны, это просто 2D-данные, к которым я хочу получить доступ.

Основная вещь, которую я пытаюсь сделать, это вычислить два варианта 2D-данных (в данном случае (6,60) ), а затем эффективно выбрать те, у которых наименьшая общая сумма — вот откуда берется двоичный (6,) массив.

Пример: для B = [1,0,1,0,1,0] того, что я хочу получить, равно стекированию

 A[1,0,:]
A[0,1,:]
A[1,2,:]
A[0,3,:]
A[1,4,:]
A[0,5,:]
  

но я хотел бы сделать это путем прямого индексирования, а не цикла for.

Я пробовал A[B], A[:,B,:], A[B,:,:] A[:,:,B] , но ни один из них не предоставил желаемую (6,60) матрицу.

 import numpy as np
A = np.array([[4, 4, 4, 4, 4, 4], [1, 1, 1, 1, 1, 1]])
A = np.atleast_3d(A)
A = np.tile(A, (1,1,60)
B = np.array([1, 0, 1, 0, 1, 0])
A[B]
  

Ожидаемые результаты представляют собой (6,60) массив, содержащий элементы из A, как описано выше, полученный либо (2,6,60) или (6,6,60) .

Заранее благодарю вас, Линус

Вы можете сгенерировать диапазон индексов, которые вы хотите перебрать, в вашем случае от 0 до 5:

 count = A.shape[1]

indices = np.arange(count)  # np.arange(6) for your particular case

>>> print(indices)
array([0, 1, 2, 3, 4, 5])
  

И затем вы можете использовать это для расширенной индексации:

 result_array = A[B[indices], indices, :]
  

Если вы всегда используете полный диапазон от 0 до длины — 1 (т. Е. От 0 до 5 в вашем случае) второй оси A в порядке возрастания, вы можете упростить это до:

 result_array = A[B, indices, :]
# or the ugly result_array = A[B, np.arange(A.shape[1]), :]
  

Или даже это, если это всегда 6:

 result_array = A[B, np.arange(6), :]
  

Альтернативное решение с использованием np.take_along_axis (из версии 1.15 — docs)

 import numpy as np
x = np.arange(2*6*6).reshape((2,6,6))
m = np.zeros(6, int)
m[0] = 1
#example: [1, 0, 0, 0, 0, 0]

np.take_along_axis(x, m[None, :, None], 0)   #add dimensions to mask to match array dimensions

>>array([[[36, 37, 38, 39, 40, 41],
        [ 6,  7,  8,  9, 10, 11],
        [12, 13, 14, 15, 16, 17],
        [18, 19, 20, 21, 22, 23],
        [24, 25, 26, 27, 28, 29],
        [30, 31, 32, 33, 34, 35]]])