pandas - создает новый столбец на основе значения строки 'next'

Question

У меня есть следующий фрейм данных:

   date      country   
   6/1/18    USA
   6/1/18    BEL
   6/4/18    USA
   6/5/18    BEL
   6/6/18    USA

Я хочу создать столбец, в котором будет указана следующая дата для соответствующей страны. Вы можете предположить, что даты упорядочены, если это поможет. Если дата является последней датой для страны, вы можете заполнить следующую дату той же датой или нулевым значением.

   date      country   next_date
   6/1/18    USA       6/4/18
   6/1/18    BEL       6/5/18
   6/4/18    USA       6/6/18
   6/5/18    BEL       6/5/18
   6/6/18    USA       6/6/18

Answer 1

Вы также можете использовать transform. Я решил оставить окончательные даты в виде значений NaN.

>>> df.assign(next_date=df.groupby('country')['date'].transform(
                  lambda group: group.shift(-1)))
     date country next_date
0  6/1/18     USA    6/4/18
1  6/1/18     BEL    6/5/18
2  6/4/18     USA    6/6/18
3  6/5/18     BEL       NaN
4  6/6/18     USA       NaN

Answer 2

Вы можете использовать groupby и shift для заполнения следующего доступного значения, однако в этом случае строки без следующей доступной даты останутся как NaN:

df.assign(ndate=df.groupby('country').date.shift(-1))

     date country   ndate
0  6/1/18     USA  6/4/18
1  6/1/18     BEL  6/5/18
2  6/4/18     USA  6/6/18
3  6/5/18     BEL     NaN
4  6/6/18     USA     NaN

Если вы хотите заполнить эти значения последней датой, вы можете просто fillna с помощью столбца date:

df.assign(ndate=df.groupby('country').date.shift(-1)).fillna({'ndate': df.date})

     date country   ndate
0  6/1/18     USA  6/4/18
1  6/1/18     BEL  6/5/18
2  6/4/18     USA  6/6/18
3  6/5/18     BEL  6/5/18
4  6/6/18     USA  6/6/18