yobogoya
yobogoya

Reputation: 594

(pandas) Create new column based on first element in groupby object

Say I have the following dataframe:

>>> df = pd.DataFrame({'Person': ['bob', 'jim', 'joe', 'bob', 'jim', 'joe'], 'Color':['blue', 'green', 'orange', 'yellow', 'pink', 'purple']})
>>> df

    Color Person
0    blue    bob
1   green    jim
2  orange    joe
3  yellow    bob
4    pink    jim
5  purple    joe

And I want to create a new column that represents the first color seen for each person:

     Color Person First Color
0    blue    bob        blue
1   green    jim       green
2  orange    joe      orange
3  yellow    bob        blue
4    pink    jim       green
5  purple    joe      orange

I have come to a solution but it seems really inefficient:

>>> df['First Color'] = 0
>>> groups = df.groupby(['Person'])['Color']
>>> for g in groups:
...    first_color = g[1].iloc[0]
...    df['First Color'].loc[df['Person']==g[0]] = first_color

Is there a faster way to do this all at once where it doesn't have to iterate through the groupby object?

Upvotes: 4

Views: 1517

Answers (2)

MaxU - stand with Ukraine
MaxU - stand with Ukraine

Reputation: 210832

use transform() method:

In [177]: df['First_Col'] = df.groupby('Person')['Color'].transform('first')

In [178]: df
Out[178]:
    Color Person First_Col
0    blue    bob      blue
1   green    jim     green
2  orange    joe    orange
3  yellow    bob      blue
4    pink    jim     green
5  purple    joe    orange

Upvotes: 5

jezrael
jezrael

Reputation: 862471

You need transform with first:

print (df.groupby('Person')['Color'].transform('first'))
0      blue
1     green
2    orange
3      blue
4     green
5    orange
Name: Color, dtype: object

df['First_Col'] = df.groupby('Person')['Color'].transform('first')
print (df)
    Color Person First_Col
0    blue    bob      blue
1   green    jim     green
2  orange    joe    orange
3  yellow    bob      blue
4    pink    jim     green
5  purple    joe    orange

Upvotes: 10

Related Questions