Доступ к Numpy Arary на основе других массивов

Question

У меня есть массив N размером [2, 10]. У меня есть другой массив B длиной 10, значения которого находятся в диапазоне от 0 до 1. Теперь я хочу получить доступ к A, беря для каждого столбца строку, указанную B, что означает: Выход должен иметь длину 10, i-тая запись в результирующем массиве будет A [0, i], если B [i] == 0, иначе A [1, i]. Как мне индексировать это правильно / наиболее элегантно? Спасибо.

Answer 1

Вы можете использовать numpy.choose для этого:

>>> a
array([[ -3,   2,   5,   0,  -2,  12,   5,  -1, -12,   5],
       [ -8, -12,  15,   6,   2,  -3,   4, -17,   0,   7]])
>>> b
array([1, 0, 0, 1, 1, 1, 1, 1, 0, 1])
>>> np.choose(b, a)
array([ -8,   2,   5,   6,   2,  -3,   4, -17, -12,   7])

Здесь первый параметр b - это массив, который содержит значения вариантов, а следующий параметр - это 2d-массив значений для выбора.

Как видите, для индексов i, где b[i] == 0, затем соответствующее значение первой строки a[0][i], а для b[i] == 1 мы получаем a[1][i]. Это может быть легко обобщено для большего количества значений.

или мы можем использовать numpy.where и передавать подмассивы:

>>> np.where(b, a[1], a[0])
array([ -8,   2,   5,   6,   2,  -3,   4, -17, -12,   7]) 

Answer 2

Я считаю, что это будет работать:

A[1][i] if B[i] == 0 else A[0][i]