Dataframe Pandas - Calculate column based on other

Question

I have a dataframe in the following format:

colunas = [
    'COMEDY',
    'CRIME',
    'Classe Prevista'
]

precisao_df = pd.DataFrame(columns=colunas)
precisao_df['COMEDY'] = y_pred_proba[:,0]
precisao_df['CRIME'] = y_pred_proba[:,1]
precisao_df

And I want to do the following: for each line, if comedy > crime, expected class = comedy, if not crime. That is, the predicted class is the one my model gave the most probability of being. How can I make this expression?

Ex:

    COMEDY    CRIME      Classe prevista
0    0.9       0.1          COMEDY
1    0.42      0.58         CRIME

Answer 1

You can use the function where the package Numpy, just inform a condition and the values to be used if true or false.

df['Classe Prevista'] = np.where(df.COMEDY > df.CRIME, 'COMEDY', 'CRIME')
#                                                         |        |
#                                                         - True   |
#                                                                  -  False

Exit:

     COMEDY     CRIME Classe Prevista
0  0.909657  0.090343          COMEDY
1  0.180432  0.819568           CRIME
2  0.756018  0.243982          COMEDY

See working in repl it.

Answer 2

Create a function that determines the class you want to place in the column:

def acha_classe(registro):
    if registro['COMEDY'] > registro['CRIME']:
        return 'COMEDY'
    elif registro['COMEDY'] < registro['CRIME']:
        return 'CRIME'
    else:
        return 'UNKNOWN'

Just apply (apply) that function:

df['Classe Prevista'] = df.apply(acha_classe, axis=1)