Home > Software design >  Remove element from a list based on condition in pandas dataframe
Remove element from a list based on condition in pandas dataframe

Time:04-30

a= {'A' : [1, 2,3,4],
'B' : ['FOOTBALL','BASKETBALL','HANDBALL','VOLLEYBALL'],
'C' : [[5,10,15,40],[1,4],[20,10,40],[10,40]]
}

How can I remove the element 40 from C if B is different to FOOTBALL Like this :

    A   B   C
0   1   FOOTBALL    [5, 10, 15, 40]
1   2   BASKETBALL  [1, 4]
2   3   HANDBALL    [20, 10]
3   4   VOLLEYBALL  [10]

CodePudding user response:

You can filter both sides in mask and for remove 40 in list use lambda function:

df = pd.DataFrame(a)

m = df['B'].ne('FOOTBALL')
df.loc[m, 'C'] = df.loc[m, 'C'].apply(lambda x: [y for y in x if y!=40])
print (df)
   A           B                C
0  1    FOOTBALL  [5, 10, 15, 40]
1  2  BASKETBALL           [1, 4]
2  3    HANDBALL         [20, 10]
3  4  VOLLEYBALL             [10]

Alternative:

m = df['B'].ne('FOOTBALL')
df.loc[m, 'C'] = [[y for y in x if y!=40] for x in df.loc[m, 'C']]
  • Related