Как разделить массив с помощью двоичного списка?

Если у меня есть NumPy массив сказать

A = [[1,2],[3,4],[5,6],[7,8]]

и список

L = [1,0,1,1]

Есть ли способ разделить А по оси 0, если они составляют 1/0 в L? Это был бы мой желаемый результат:

A1 = [[1,2],[5,6],[7,8]]
A2 = [[3,4]]

1 ответ

Решение

Поскольку L является двоичным, вы можете конвертировать L для логического типа, а затем использовать логическое индексирование:

A = np.array([[1,2],[3,4],[5,6],[7,8]])
L = np.array([1,0,1,1])

L = L.astype(bool)
A1, A2 = A[L], A[~L]

A1
#array([[1, 2],
#       [5, 6],
#       [7, 8]])

A2
#array([[3, 4]])
Другие вопросы по тегам