Groupby + conditional from another column to create new one

问题

I am trying to capture the date of the “visit_num==2” of “users” in a new column ("2nd_visit_date")

Here's the code (including the new column I want to create)

df=pd.DataFrame({'user':[1,1,2,2,2,3,3,3,3,3,4,4],
                  'date':['1995-09-01','1995-09-02','1995-10-03','1995-10-04','1995-10-05','1995-11-07','1995-11-08','1995-11-09','1995-11-10','1995-11-15','1995-12-18','1995-12-20'],
                  'visit_num':[1,2,1,2,3,1,2,3,4,5,1,2],
                  '2nd_visit_date':['1995-09-02','1995-09-02','1995-10-04','1995-10-04','1995-10-04','1995-11-08','1995-11-08','1995-11-08','1995-11-08','1995-11-08','1995-12-20','1995-12-20']})

So I get:

user    date    visit_num   2nd_visit_date
 1   1995-09-01     1        1995-09-02
 1   1995-09-02     2        1995-09-02
 2   1995-10-03     1        1995-10-04
 2   1995-10-04     2        1995-10-04
 2   1995-10-05     3        1995-10-04
 3   1995-11-07     1        1995-11-08
 3   1995-11-08     2        1995-11-08
 3   1995-11-09     3        1995-11-08
 3   1995-11-10     4        1995-11-08
 3   1995-11-15     5        1995-11-08
 4   1995-12-18     1        1995-12-20
 4   1995-12-20     2        1995-12-20

I tried the following code, but it did not work:

df["2nd_visit_date"] = df.groupby("user")["date"].transform(df['visit_num']==2)

Any help will be very much appreciated. Thanks.

回答1:

Let say this is your original df:

df

   user    date    visit_num
0   1   1995-09-01  1
1   1   1995-09-02  2
2   2   1995-10-03  1
3   2   1995-10-04  2
4   2   1995-10-05  3
5   3   1995-11-07  1
6   3   1995-11-08  2
7   3   1995-11-09  3
8   3   1995-11-10  4
9   3   1995-11-15  5
10  4   1995-12-18  1
11  4   1995-12-20  2

You can first create a dataframe for second visits (and change column name):

df_2 = df[df.visit_num==2][['user', 'date']]
df_2.columns = ['user', '2nd_visit_date']
df_2


   user 2nd_visit_date
1   1   1995-09-02
3   2   1995-10-04
6   3   1995-11-08
11  4   1995-12-20

And merge it with your original df

pd.merge(df, df_2, on='user', how='left')

    user    date    visit_num   2nd_visit_date
0   1   1995-09-01      1         1995-09-02
1   1   1995-09-02      2         1995-09-02
2   2   1995-10-03      1         1995-10-04
3   2   1995-10-04      2         1995-10-04
4   2   1995-10-05      3         1995-10-04
5   3   1995-11-07      1         1995-11-08
6   3   1995-11-08      2         1995-11-08
7   3   1995-11-09      3         1995-11-08
8   3   1995-11-10      4         1995-11-08
9   3   1995-11-15      5         1995-11-08
10  4   1995-12-18      1         1995-12-20
11  4   1995-12-20      2         1995-12-20

来源：https://stackoverflow.com/questions/60870324/groupby-conditional-from-another-column-to-create-new-one

标签

python

conditional-statements

pandas-groupby