CODEWITHSUNDEEP
pythonpandas
spitfiredd
spitfiredd

Reputation: 3135

Pandas: Flatten column based on condition?

I am trying to flatten rows and keep the info from the rows I want.

What I have:

id  var1  var2 var3
1      Y     N    Y
1      N          Y
2      Y          N
2      N     Y    N
2      Y     N    Y

What I would like:

id  var1  var2 var3
1      Y     N    Y
2      Y     Y    Y

Essentially, it would check if there is a Y/N and always give priority to a Y. Also there are more columns than var1, var2, var3; so I would like something more general so I could apply to other columns as well.

Upvotes: 4

Views: 471

Answers (3)

jezrael
jezrael

Reputation: 862901

You can use replace + groupby + GroupBy.max + replace + reset_index:

df1 = df.replace({'Y':1,'N':0, np.nan:-1})
        .groupby('id')
        .max()
        .replace({1:'Y', 0:'N',-1:np.nan})
        .reset_index()
print (df1)

   id var1 var2 var3
0   1    Y    N    Y
1   2    Y    Y    Y

EDIT:

df = pd.DataFrame({
'id': [1, 1, 2, 2, 3, 3], 
'var2': ['N', 'N', 'N', 'Y', 'N', np.nan], 
'var1': ['Y', 'Y', 'Y', 'N', 'Y', np.nan], 
'var3': [np.nan, np.nan, np.nan, 'N', np.nan, 'Y']
})

print (df)
   id var1 var2 var3
0   1    Y    N  NaN
1   1    Y    N  NaN
2   2    Y    N  NaN
3   2    N    Y    N
4   3    Y    N  NaN
5   3  NaN  NaN    Y

You can create another dict dynamically:

#check all unique values without column id
print (df.set_index('id').stack(dropna=False).unique())
['Y' 'N' nan]

#create dict for first replace
d = {'Y':1,'N':0, np.nan:-1}
#swap keys, values in dict for another replace
d1 = {v: k for k, v in d.items()}

df1 = df.replace(d).groupby('id').max().replace(d1).reset_index()
print (df1)
   id var1 var2 var3
0   1    Y    N  NaN
1   2    Y    Y    N
2   3    Y    N    Y

EDIT1:

Solution if only Y,N and NaN in var1 - varN columns:

varNAN = 'A'
print (df.fillna(varNAN).groupby('id').max().replace({varNAN:np.nan}).reset_index())
   id var1 var2 var3
0   1    Y    N  NaN
1   2    Y    Y    N
2   3    Y    N    Y

Upvotes: 3

Allen Qin
Allen Qin

Reputation: 19957

If there's only 'Y', 'N' and NAN in your dataframe, this is a simpler way to do it.

Setup

df = pd.DataFrame({'id': {0: 1, 1: 1, 2: 2, 3: 2, 4: 2},
 'var1': {0: 'Y', 1: 'N', 2: 'Y', 3: 'N', 4: 'Y'},
 'var2': {0: 'N', 1: np.nan, 2: np.nan, 3: 'Y', 4: 'N'},
 'var3': {0: 'Y', 1: 'Y', 2: 'N', 3: 'N', 4: 'Y'}})

Out[45]: 
   id var1 var2 var3
0   1    Y    N    Y
1   1    N  NaN    Y
2   2    Y  NaN    N
3   2    N    Y    N
4   2    Y    N    Y

Solution

#Fill na with 'A' and get the max (Y>N>A) from each column.
df.fillna('A').groupby('id').max().reset_index()
Out[46]: 
   var1 var2 var3
id               
1     Y    N    Y
2     Y    Y    Y

Upvotes: 1

Scott Boston
Scott Boston

Reputation: 153460

Let's try, you can use groupby and sum to act like an OR, hence "giving Y priority":

df1 = df.replace({'Y':True,'N':False})

df_out = (df1.groupby('id').sum(skipna=False)
         .astype(bool)
         .replace({True:'Y',False:'N'})
         .reset_index())

print(df_out)

Output:

   id var1 var2 var3
0   1    Y    N    Y
1   2    Y    Y    Y

Upvotes: 5

Related Questions