Как я могу создать новый столбец pandas на основе максимального значения других столбцов pandas и связанного столбца строки?

Question

У меня есть следующий фрейм данных:

id    pred_1    conf_1   pred_2    conf_2   pred_3    conf_3
1     dog         .5       cat       .7      fish       .3
2     cat         .1       fish      .8      dog        .4
3     fish        .7       pig       .3      cat        .6

Я хочу создать новый столбец most likely, который генерирует прогноз с наивысшим уровнем достоверности. Ie: для строки 1 наибольшее значение conf_2, поэтому значение столбца most_likely - «кошка».

Я пробовал:

def highest_confidence(row):
  if max(row.p1_conf, row.p2_conf, row.p3_conf) == row.p1_conf:
    row["most_likely"] = row.p1
elif max(row.p1_conf, row.p2_conf, row.p3_conf) == row.p2_conf:
    row["most_likely"] = row.p2
else:
    row["most_likely"] = row.p3


images.apply(highest_confidence, axis=1)

, но получил ошибку: IndexError: index 13 is out of bounds for axis 0 with size 13

Answer 1

Учитывая такой фрейм данных, где пары прогноз / достоверность находятся рядом друг с другом,

import pandas as pd
import numpy as np
import random

values = ["dog", "cat", "fish", "pig"]

df = pd.DataFrame(data={
    "pred_1": random.sample(values, 4),
    "conf_1": np.random.random(4),
    "pred_2": random.sample(values, 4),
    "conf_2": np.random.random(4),
    "pred_3": random.sample(values, 4),
    "conf_3": np.random.random(4)    
})

Мы можем создать из него временный фрейм данных, который будет просто True / False в зависимости от от того, является ли значение ячейки максимальным значением для строки:

bool_matrix = df.eq(df.max(axis=1), axis="index")

Затем мы можем сдвинуть исходный фрейм данных на один столбец вправо, чтобы соответствовать логическим значениям. Обратите внимание, что по умолчанию сдвиг происходит в следующем столбце того же типа, поэтому мы должны преобразовать числа в строки:

shifted = df.astype(str).shift(periods=1, axis="columns")

И, наконец, мы go по каждому столбцу, фильтруем и объединяем:

hits = []
for column in shifted.columns:
    hits.append(shifted.loc[bool_matrix[column], column])

df["most_likely"] = pd.concat(hits)

Так как мы избегаем перебора строк, это должно быть достаточно быстро.