在pandas dataframe列表推导式中的If ElseIf Else条件

Question

我有一个包含11列的数据框：Status1-Status5，Time1-Time5和Time_Min。

df = pd.DataFrame([[100,200,150,400,500,'a','b','a','c','a',100], [300,400,200,500,250,'b','b','c','c','c',200]], columns=['TIME_1', 'TIME_2', 'TIME_3', 'TIME_4', 'TIME_5','STATUS_1','STATUS_2','STATUS_3','STATUS_4','STATUS_5','TIME_MIN'])

我希望能够复制我目前在SAS中拥有的代码，该代码执行以下操作：

IF TIME_1 = TIME_MIN THEN STATUS = STATUS_1;
ELSE IF TIME_2 = TIME_MIN THEN STATUS = STATUS_2;
ELSE IF TIME_3 = TIME_MIN THEN STATUS = STATUS_3;
ELSE IF TIME_4 = TIME_MIN THEN STATUS = STATUS_4;
ELSE STATUS = STATUS_5;

列 STATUS 的预期输出应为

['a','c']

我尝试构建类似的东西（需要使用else if进行扩展）

df['STATUS'] = [a if x == y else b for x,y,a,b in df[['TIME_MIN','TIME_1','STATUS_1','STATUS_2']]]

但是这只会导致错误。我确定这是一个简单的修复，但我无法完全弄清楚。

- DarknessFalls

---

你尝试过使用 np.where 吗？ - rpanai

那会产生哪个错误？ - Nico Albers

@Nico Albera - ValueError: 分配的值过多，无法拆包。 - DarknessFalls

欢迎来到StackOverflow！请阅读有关如何提问（特别是如何创建良好的示例）的内容，以获得良好的回复。提供一个示例DataFrame和期望的输出将有助于更好地回答您的问题。 - Alex

@Alex 示例数据已添加 - DarknessFalls

Answer 1

您可以编写一个函数。

def get_status(df):
    if df['TIME_1'] == df['TIME_MIN']:
        return df['STATUS_1']
    elif df['TIME_2'] == df['TIME_MIN']:
        return df['STATUS_2']
    elif df['TIME_3'] == df['TIME_MIN']:
        return df['STATUS_3']
    elif df['TIME_4'] == df['TIME_MIN']:
        return df['STATUS_4']
    else:
        return df['STATUS_5']

df['STATUS'] = df.apply(get_status, axis = 1)

或者使用一个非常嵌套的 np.where。

df['STATUS'] = np.where(df['TIME_1'] == df['TIME_MIN'], df['STATUS_1'],\ 
        np.where(df['TIME_2'] == df['TIME_MIN'], df['STATUS_2'],\
        np.where(df['TIME_3'] == df['TIME_MIN'], df['STATUS_3'],\
        np.where(df['TIME_4'] == df['TIME_MIN'], df['STATUS_4'], df['STATUS_5']))))

- Vaishali

---

谢谢。我已经让函数选项正常工作了，但希望能够创建一些更简洁的东西。特别是因为将来可能会有超过5个选项。 - DarknessFalls

Answer 2

虽然不太美观，但你可以使用等式广播和.eq方法来实现。

m = df.iloc[:, :5].eq(df['TIME_MIN'], axis=0)
m.columns = 'STATUS_' + m.columns.str.extract('TIME_(.*)')
df['STATUS'] = df[m].bfill(axis=1).iloc[:, 0]

- Alex

Answer 3

您可以使用条件和选择

df = pd.DataFrame([[100,200,150,400,500,'a','b','a','c','a',100], [300,400,200,500,250,'b','b','c','c','c',200]], columns=['TIME_1', 'TIME_2', 'TIME_3', 'TIME_4', 'TIME_5','STATUS_1','STATUS_2','STATUS_3','STATUS_4','STATUS_5','TIME_MIN'])


condition= [df['TIME_1'] == df['TIME_MIN'],
            df['TIME_2'] == df['TIME_MIN'],
            df['TIME_3'] == df['TIME_MIN'],
            df['TIME_4'] == df['TIME_MIN'],
            df['TIME_4'] == df['TIME_MIN']]

choice= [df['STATUS_1'],df['STATUS_2'],df['STATUS_3'],df['STATUS_4'],df['STATUS_5']]

df['STATUS'] =np.select(condition,choice,default="")

col_required=['TIME_1','TIME_2','TIME_3','TIME_4','TIME_5','TIME_MIN','STATUS']
df=df[col_required]
df

输出

    TIME_1  TIME_2  TIME_3  TIME_4  TIME_5  TIME_MIN    STATUS
0   100 200 150 400 500 100 a
1   300 400 200 500 250 200 c

- Vishvas Chauhan