Матрица NumPy 2D: возвращает истинные индексы строк, где столбец истинен

Question

У меня есть два массива в NumPy:

a1 = np.array([[0, 0, 0],
               [1, 1, 0],
               [0, 1, 0],
               [1, 1, 1]])

a2 = np.array([1, 2])

, и мне нужно знать индексы строк a1, где оба столбца 1 и 2 (заданные в a2) равны 1.

В приведенном выше примере это будет только последняя строка с индексом 3. Не имеет значения, если другие столбцы равны 1.

Может ли кто-нибудь помочь мне сделать это без сильных циклов?

Answer 1

Использование простых индексов и np.all даст вам строки, где все элементы равны 1:

np.all(a1[:, a2], axis=1)

Это даст вам маску строки, соответствующую этому критерию. Как это:

array([False, False, False,  True])

Затем, следуя @ rahlf23 комментарию, вы можете использовать np.where, чтобы найти индексы, соответствующие следующим критериям:

np.where(np.all(a1[:, a2], axis=1))[0]

Выход :

array([3])