вычислить разницу дат между минимумом в том же столбце в зависимости от состояния другого столбца, используя Python

Question

d = {'country': ['US', 'US', 'US','US', 'US', 'US', 'UK', 'UK','UK','UK','UK'], 
     'status': [0, 0, 0, 0, 1,1,0, 0,  0, 1,1],
     'count':[0, 1, 10, 20,30,40,0,1,2,4,6],
     'date':['2020-04-05', '2020-04-06', '2020-04-07', '2020-04-11', '2020-04-12', 
             '2020-04-13', '2020-04-02', '2020-04-03', '2020-04-05', '2020-04-06', '2020-04-07']}
df = pd.DataFrame(data=d)

Первое, что я подумал, это использовать

df['day_diff']=df.groupby(['country', 'status'])['date'].apply(lambda x: x - x.min())

Однако 0 дней для US status0 должно появиться 06.04.2020, а не 2020-04-05, в котором count равно 0 (извините, я не знаю, как правильно разместить здесь вывод). Таким образом, должно быть два условия для группы, применяемой к возвращаемому day_diff

1. count> 0

2. date - date.min ()

   country status  count   date    day_diff
       US  0   0   2020-04-05  0 days
   US  0   1   2020-04-06  1 days
   US  0   10  2020-04-07  2 days
   

   3 США 0 20 2020-04-11 6 дней 4 США 1 30 2020-04-12 0 дней 5 США 1 40 2020-04-13 1 дней 6 Великобритания 0 0 2020-04-02 0 дней 7 Великобритания 0 1 2020 -04-03 1 день 8 Великобритания 0 2 2020-04-05 3 дня 9 Великобритания 1 4 2020-04-06 0 дней 10 Великобритания 1 6 2020-04-07 1 день

Надеюсь получить результат, подобный этому (конечно не работает):

df.groupby(['country', 'status'])['date'].apply(lambda x: x - x.min() if df['count']>0)

Чтобы вычислить разницу между датой и самой ранней датой в группе, в зависимости от того, подсчитали ли первая дата выше 0.

Заранее спасибо.

Answer 1

Вы можете использовать groupby().transform() для извлечения минимальной даты:

min_dates = (df.date.mask(df['count'].eq(0))
               .groupby([df['country'],df['status']])
               .transform('min')
            )

df['date_diff'] = df['date'] - min_dates

Вывод:

   country  status  count       date date_diff
0       US       0      0 2020-04-05   -1 days
1       US       0      1 2020-04-06    0 days
2       US       0     10 2020-04-07    1 days
3       US       0     20 2020-04-11    5 days
4       US       1     30 2020-04-12    0 days
5       US       1     40 2020-04-13    1 days
6       UK       0      0 2020-04-02   -1 days
7       UK       0      1 2020-04-03    0 days
8       UK       0      2 2020-04-05    2 days
9       UK       1      4 2020-04-06    0 days
10      UK       1      6 2020-04-07    1 days