CODEWITHSUNDEEP
python-3.xpandasnumpywhere-clausepandas-groupby
Rachel
Rachel

Reputation: 1967

Create conditional column in pandas

I am trying to create a conditional column in pandas. Here is what the dataframe looks like.

    data = [{"owner" : "john", "dog" : 'magie', "dog_is_fluffy" : 1},
            {"owner" : "john", "dog" : 'stellar', "dog_is_fluffy" : 0}, 
            {"owner" : "lisa", "dog" : 'mollie' , "dog_is_fluffy" : 0},
            {"owner" : "lisa", "dog" : 'rex', "dog_is_fluffy" : 0},
            {"owner" : "john", "dog" : 'luns', "dog_is_fluffy" : 1}]

    df = pd.DataFrame(data)

As you can see, my data shows dogs and their owners. We also know if the dog is fluffy. I want to create two columns fluffy_dogs_owned and owner_has_fluffy_dog.

The result I am looking for is:

data_result = [{"owner" : "john", "dog" : 'magie', "dog_is_fluffy" : 1, "fluffy_dogs_owned" : 2, "owner_has_fluffy_dog" : 1},
        {"owner" : "john", "dog" : 'stellar', "dog_is_fluffy" : 0, "fluffy_dogs_owned" : 2, "owner_has_fluffy_dog" : 1}, 
        {"owner" : "lisa", "dog" : 'mollie' , "dog_is_fluffy" : 0, "fluffy_dogs_owned" : 0, "owner_has_fluffy_dog" : 0},
        {"owner" : "lisa", "dog" : 'rex', "dog_is_fluffy" : 0, "fluffy_dogs_owned" : 0, "owner_has_fluffy_dog" : 0},
        {"owner" : "john", "dog" : 'luns', "dog_is_fluffy" : 1, "fluffy_dogs_owned" : 2, "owner_has_fluffy_dog" : 1}]

df_result = pd.DataFrame(data_result)

I thought about using df.groupby() and np.where but I can't make it work so far. Any ideas?

Upvotes: 1

Views: 78

Answers (1)

jezrael
jezrael

Reputation: 862791

Use GroupBy.transform for return Series with same size like original Dataframe with sum and then compare column for not equal by Series.ne with casting to integer

df['fluffy_dogs_owned'] = df.groupby('owner')['dog_is_fluffy'].transform('sum')
df['owner_has_fluffy_dog'] = df['fluffy_dogs_owned'].ne(0).astype(int)

Or with Series.clip:

df['owner_has_fluffy_dog'] = df['fluffy_dogs_owned'].clip(upper=1)

print (df)
       dog  dog_is_fluffy owner  fluffy_dogs_owned  owner_has_fluffy_dog
0    magie              1  john                  2                     1
1  stellar              0  john                  2                     1
2   mollie              0  lisa                  0                     0
3      rex              0  lisa                  0                     0
4     luns              1  john                  2                     1

Upvotes: 2

Related Questions