Как найти индекс столбца значения данной строки в numpy?

Question

Я хочу, чтобы выполнить что-то относительно простой в NumPy:

• Если есть один в строке, возвращает индекс столбца, содержащего одну +1.
• Если есть ноль или несколько один в строке возвращения 0.

Однако я в конечном итоге с довольно сложным кодом:

predictions = np.array([[1,-1,-1,-1],[-1,1,1,-1],[-1,-1,-1,1],[-1,-1,-1,-1]]) 

one_count = (predictions == 1).sum(1) 
valid_rows_idx = np.where(one_count==1) 

result = np.zeros(predictions.shape[0]) 
for idx in valid_rows_idx: 
    result[idx] = np.where(predictions[idx,:]==1)[1] + 1 

Если я печатаю result, программа печатает [ 1. 0. 4. 0.], которые является желаемым результатом.

Мне было интересно, существует ли более простой способ записи последней строки с использованием numpy.

Answer 1

Я не уверен, что это лучше или нет, но вы можете попробовать использовать argmax. Кроме того, вам не нужно использовать цикл и np.where, чтобы получить действительные показатели:

predictions = np.array([[1,-1,-1,-1],[-1,1,1,-1],[-1,-1,-1,1],[-1,-1,-1,-1]]) 

idx = (predictions == 1).sum(1) == 1 
result = np.zeros(predictions.shape[0]) 
result[idx] = (predictions[idx]==1).argmax(axis=1) + 1 

In [55]: result 
Out[55]: array([ 1., 0., 4., 0.]) 

Или вы могли бы сделать все эти вещи с одной линии с использованием np.where и argmax:

predictions = np.array([[1,-1,-1,-1],[-1,1,1,-1],[-1,-1,-1,1],[-1,-1,-1,-1]]) 

In [72]: np.where((predictions==1).sum(1)==1, (predictions==1).argmax(axis=1)+1, 0) 
Out[72]: array([1, 0, 4, 0])