如何根据其他列的条件分配新列？

Question

我正在尝试根据 A、B、C、D 中存在的值添加一列“flag_column”。

即如果 A/B/C/DI 中有一个值想创建一个新列，'flag' 表示包含一个值的列名。

  A B C D counts flag
0 1 0 0 0  1     A
1 0 1 0 0  1     B
2 1 0 0 0  1     A
3 0 0 1 0  1     C
4 0 1 0 0  1     B

注意：只有一列（A 到 D）包含值，因此计数始终为 1。

我试过了：

if [df['A'] == 1] == True:
    df['flag'] = 'A'
elif [df['B'] == 1] == True:
    df['flag'] = 'B'
elif [df['C'] == 1] == True:
    df['flag'] = 'C'  
else:
    df['flag'] = 'D'    

我也试过：

df['flag'] = np.where(df['A'] == 1, 'A', False)
df['flag'] = np.where(df['B'] == 1, 'B', False)
df['flag'] = np.where(df['C'] == 1, 'C', False)
df['flag'] = np.where(df['D'] == 1, 'D', False)

我也尝试过迭代地遍历每个“类别”并分配一个标志值，但是在这些情况下它也会覆盖。

如果有一种方法可以迭代地执行此操作，那将是理想的。但是，对这个（简单）问题的任何帮助将不胜感激！

Answer 1

我们可以idxmax在axis=1这里使用：

df['flag'] = df.loc[:, 'A':'D'].idxmax(axis=1)

   A  B  C  D flag
0  1  0  0  0    A
1  0  1  0  0    B
2  1  0  0  0    A
3  0  0  1  0    C
4  0  1  0  0    B