How to separate in x and y being all but the last x and the last y?

Question

I am doing a function to generalize a train_test_split, and for that desire separates within the function the X and the y. Currently my code is like this:

python
def split(dados):
    SEED = 42367
    X = dados[["NU_NOTA_MT", "NU_NOTA_LC"]]
    y = dados["nota_total"]
    train_x, test_x, train_y, test_y = train_test_split(X, y, random_state=SEED)
    return train_x, test_x, train_y, test_y

Note that when I define X and y I need to pass the name of the variables inside the function... It’s like I’m just passing the data he understands that all -1 is X and the last one is y?

Answer 1

You can pass the column names as a list and slice them:

def split(dados):
    SEED = 42367
    X = dados[dados.columns.to_list()[:-1]]
    y = dados[dados.columns.to_list()[-1:]]
    train_x, test_x, train_y, test_y = train_test_split(X, y, random_state=SEED)
    return train_x, test_x, train_y, test_y