Numpy: ¿Cómo encontrar el primer valor distinto de cero en cada columna de una matriz numpy?

Supongamos que tengo una matriz numpy de la forma:

arr=numpy.array([[1,1,0],[1,1,0],[0,0,1],[0,0,0]]) 

Quiero encontrar los índices del primer índice (para cada columna) donde el valor es distinto de cero.

Así que en este caso, me gustaría que se devuelva lo siguiente:

 [0,0,2] 

¿Cómo hago para esto?

Índices de primeras apariciones.

Use np.argmax largo de ese eje (eje cero para las columnas aquí) en la máscara de los no ceros para obtener los índices de las primeras matches (valores verdaderos) –

 (arr!=0).argmax(axis=0) 

Extendiéndonos para cubrir el especificador genérico de ejes y para los casos en que no se encuentren no ceros a lo largo de ese eje para un elemento, tendríamos una implementación como tal:

 def first_nonzero(arr, axis, invalid_val=-1): mask = arr!=0 return np.where(mask.any(axis=axis), mask.argmax(axis=axis), invalid_val) 

Tenga en cuenta que dado que argmax() en todos los valores False devuelve 0 , por lo tanto, si el valor de invalid_val necesario es 0 , tendríamos la salida final directamente con mask.argmax(axis=axis) .

Ejecuciones de muestra

 In [296]: arr # Different from given sample for variety Out[296]: array([[1, 0, 0], [1, 1, 0], [0, 1, 0], [0, 0, 0]]) In [297]: first_nonzero(arr, axis=0, invalid_val=-1) Out[297]: array([ 0, 1, -1]) In [298]: first_nonzero(arr, axis=1, invalid_val=-1) Out[298]: array([ 0, 0, 1, -1]) 

Extendiéndose para cubrir todas las operaciones de comparación.

Para encontrar los primeros zeros , simplemente use arr==0 como mask para usar en la función. Para los primeros iguales a un cierto valor val , use arr == val y así sucesivamente para todos los casos de comparisons posibles aquí.


Índices de últimas ocurrencias.

Para encontrar los últimos que coincidan con un cierto criterio de comparación, necesitamos voltear a lo largo de ese eje y usar la misma idea de usar argmax y luego compensar el giro desplazando la longitud del eje, como se muestra a continuación:

 def last_nonzero(arr, axis, invalid_val=-1): mask = arr!=0 val = arr.shape[axis] - np.flip(mask, axis=axis).argmax(axis=axis) - 1 return np.where(mask.any(axis=axis), val, invalid_val) 

Ejecuciones de muestra

 In [320]: arr Out[320]: array([[1, 0, 0], [1, 1, 0], [0, 1, 0], [0, 0, 0]]) In [321]: last_nonzero(arr, axis=0, invalid_val=-1) Out[321]: array([ 1, 2, -1]) In [322]: last_nonzero(arr, axis=1, invalid_val=-1) Out[322]: array([ 0, 1, 1, -1]) 

Una vez más, todos los casos de comparisons posibles aquí se cubren utilizando el comparador correspondiente para obtener la mask y luego utilizando la función enumerada.