Checking multiple columns condition in pandas

Viewed 5082

I want to create a new column in my dataframe that places the name of the column in the row if only that column has a value of 8 in the respective row, otherwise the new column's value for the row would be "NONE". For the dataframe df, the new column df["New_Column"] = ["NONE","NONE","A","NONE"]

df = pd.DataFrame({"A": [1, 2,8,3], "B": [0, 2,4,8], "C": [0, 0,7,8]})
3 Answers

Cool problem.

  1. Find the 8-fields in each row: df==8
  2. Count them: (df==8).sum(axis=1)
  3. Find the rows where the count is 1: (df==8).sum(axis=1)==1
  4. Select just those rows from the original dataframe: df[(df==8).sum(axis=1)==1]==8
  5. Find the 8-fields again: df[(df==8).sum(axis=1)==1]==8)
  6. Find the columns that hold the True values with idxmax (because True>False): (df[(df==8).sum(axis=1)==1]==8).idxmax(axis=1)
  7. Fill in the gaps with "NONE"

To summarize:

df["New_Column"] = (df[(df==8).sum(axis=1)==1]==8).idxmax(axis=1)
df["New_Column"] = df["New_Column"].fillna("NONE")
#   A  B  C New_Column
#0  1  0  0       NONE
#1  2  2  0       NONE
#2  8  4  7          A
#3  3  8  8       NONE
# I added another line as a proof of concept
#4  0  8  0          B

You can accomplish this using idxmax and a mask:

out = (df==8).idxmax(1)
m = ~(df==8).any(1) | ((df==8).sum(1) > 1)

df.assign(col=out.mask(m))

   A  B  C  col
0  1  0  0  NaN
1  2  2  0  NaN
2  8  4  7    A
3  3  8  8  NaN

Or do:

df2=df[(df==8)]
df['New_Column']=(df2[(df2!=df2.dropna(thresh=2).values[0]).all(1)].dropna(how='all')).idxmax(1)
df['New_Column'] = df['New_Column'].fillna('NONE')
print(df)

dropna + dropna again + idxmax + fillna. that's all you need for this.

Output:

   A  B  C New_Column
0  1  0  0       NONE
1  2  2  0       NONE
2  8  4  7          A
3  3  8  8       NONE
Related