熊猫：有条件地替换组中连续的行

Question

I am trying to build "episodes" from a list of transactions organized by group (patient). I used to do this with Stata, but I'm not sure how to do it in Python. In Stata, I would say something like:

by patient: replace startDate = startDate[_n-1] if startDate-endDate[_n-1]<10

In English, that meant to start with the first row of a group and check if the number of days between the startDate of that group and the endDate of the prior group was less than 10. Then, move to the next row and perform the same thing, then the next row... until you'd exhausted all rows.

I have been trying to figure out how to do the same thing in Python/Pandas and running into a wall. I could sort the dataframe by patient and date, then iterate over the entire data frame. It seems like there should be a better way to do this.

It's important that the script first compare row 2 to row 1 because, when I get to row 3, if the script has replaced the value in row 2, when I get to row 3, I want to use the replaced value, not the original value.

Sample input:

Patient    startDate    endDate  
1          1/1/2016     1/2/2016  
1          1/11/2016    1/12/2016  
1          1/28/2016    1/28/2016  
1          6/15/2016    6/16/2016  
2          3/1/2016     3/1/2016

Sample output:

Patient    startDate    endDate  
1          1/1/2016     1/2/2016  
1          1/1/2016     1/12/2016  
1          1/1/2016     1/28/2016  
1          6/15/2016    6/16/2016  
2          3/1/2016     3/1/2016

Answer 1

I think we need shift + groupby , and bfill + mask is the key

df.startDate=pd.to_datetime(df.startDate)
df.endDate=pd.to_datetime(df.endDate)

df.startDate=df.groupby('Patient').apply(lambda x : x.startDate.mask((x.startDate-x.endDate.shift(1)).fillna(0).astype('timedelta64[D]')<10).bfill()).reset_index(level=0,drop=True).fillna(df.startDate)
df
Out[495]: 
   Patient  startDate    endDate
0        1 2016-01-28 2016-01-02
1        1 2016-01-28 2016-01-12
2        1 2016-01-28 2016-01-28
3        1 2016-06-15 2016-06-16
4        2 2016-03-01 2016-03-01