Получить значение индекса из массива с условием

Question

У меня есть такой массив np.

a = [ [0. 0. 1. 0.]
  [0. 1. 0. 0.]
  [1. 0. 0. 0.]
  [0. 0. 1. 0.]
  [0. 0. 1. 0.]
]

Я хочу получить индекс всех строк в 3-м столбце, если значение элемента == 1

a[:,2:2+1]==1

В этом случаемой результат будет index = [0 3,3] Есть ли какая-нибудь функция, которую я могу использовать для этого?

Answer 1

import numpy as np
a=np.array([[0, 0, 1, 0],[0, 1, 0, 0],[1,0, 0, 0],[0, 0, 1, 0]])
index,value_first_at_index=np.where(a[:,2:3]==1)
print(index)