Quiero reconstruir la siguiente lógica con una función de transmisión numpy como np.where : desde una matriz 2d, verifique por fila si el primer elemento satisface una condición. Si la condición es verdadera, devuelva los primeros tres elementos como una fila, de lo contrario, los últimos tres elementos.
Un MWE corto en forma de bucle for que quiero eludir:
import numpy as np array = np.array([ [1, 2, 3, 4], [1, 2, 4, 2], [2, 3, 4, 6] ]) new_array = np.zeros((array.shape[0], array.shape[1]-1)) for i, row in enumerate(array): if row[0] == 1: new_array[i] = row[:3] else: new_array[i] = row[-3:]IIUC quieres algo como esto:
condition = array[:,0]==1 new_array[condition,:] = array[condition,:3] new_array[~condition,:] = array[~condition,-3:]Si quieres usar np.where :
import numpy as np array = np.array([ [1, 2, 3, 4], [1, 2, 4, 2], [2, 3, 4, 6] ]) cond = array[:, 0] == 1 np.where(cond[:, None], array[:,:3], array[:,-3:])producción:
array([[1, 2, 3], [1, 2, 4], [3, 4, 6]])EDITAR
versión un poco más concisa:
np.where(array[:, [0]] == 1, array[:,:3], array[:,-3:])