Pick for every group the next max date of a bigger group

user_18702469 0
Viewed 46

I have the following dataframe:

import pandas as pd
data = pd.DataFrame({'car_id':[1,1,1,2,2,3],'brand':['a','a','a','a','a','b'],'date': 
['01/2021','02/2022','03/2022','04/2022','05/2022','06/2022']})
data['date'] = pd.to_datetime(data['date'],format='%m/%Y')

Looks like:

'|    |   car_id | brand   | date                |
 |---:|---------:|:--------|:--------------------|
 |  0 |        1 | a       | 2021-01-01 00:00:00 |
 |  1 |        1 | a       | 2022-02-01 00:00:00 |
 |  2 |        1 | a       | 2022-03-01 00:00:00 |
 |  3 |        2 | a       | 2022-04-01 00:00:00 |
 |  4 |        2 | a       | 2022-05-01 00:00:00 |
 |  5 |        3 | b       | 2022-06-01 00:00:00 |'

I need to pick for every car_id the next available date for the same brand, taking as reference it's own car_id max(date).

car_id 1 has a max(date) of 2022-03-01, so I would need a new column after grouping where it has the date 2022-04-01, because it's the next available date (or closest if you wish) from the it's own max_date for the brand 'a' (although it belongs to a car_id 2, it's the same brand!).

car_id 2 has a max(date) of 2022-05-01, which is the max(date) also for the brand (there is no next date available), there is another max(date) in the dataset but it's from another brand! so it should be NaT.

car_id 3 has also the max_date for it's own brand, so it's same as the previous case.

The desired output should look like:

'|    |   car_id | brand   | date                |
 |---:|---------:|:--------|:--------------------|
 |  0 |        1 | a       | 2022-04-01 00:00:00 |
 |  3 |        2 | b       |        NaT          |
 |  5 |        3 | b       |        NaT          |'

I have been playing around with a simple group by max and trying to shift, but can't get it done:

data.groupby(['car_id','brand']).apply(lambda x: x['date'].max())

I can provide more context and larger examples if needed, but the pseudo code would be: for every model_id, pick its max(date) and look for the closest date in the same brand, if if max(model_id.date) == max(brand.date) then it's NaT.

-- EDIT --

Does not seem to work if you change the order on how the dataframe is created, notw for row with index 3 (car_id 2 and brand 'a') the month is now the 10th:

data = pd.DataFrame({'car_id':[1,1,1,2,2,3],'brand': 
['a','a','a','a','a','b'],'date': 
['01/2021','02/2022','03/2022','10/2022','05/2022','06/2022']})

Results in

'|    |   car_id | brand   | date    |
 |---:|---------:|:--------|:--------|
 |  0 |        1 | a       | 01/2021 |
 |  1 |        1 | a       | 02/2022 |
 |  2 |        1 | a       | 03/2022 |
 |  3 |        2 | a       | 10/2022 |
 |  4 |        2 | a       | 05/2022 |
 |  5 |        3 | b       | 06/2022 |'

Results in

car_id  brand
1       a        05/2022
2       a        10/2022
3       b            NaN
Name: date, dtype: object

Expected to car_id 1 have 10/2022 as it's the max date for the brand 'a', expected car_id to have NaN because has the max_date for the brand NaN:

car_id  brand
1       a        5/2022
2       a            NaN
3       b            NaN
Name: date, dtype: object
pandas python
1 Answers

I have had a new idea for your problem, hopefully the right one ;)

You can set up a mapping Series to find the next available date for each date per brand:

s = (data['date'].sort_values()       # ensure the dates are sorted
    .groupby(data['brand']).shift(-1) # shift to next one per group
    .set_axis(data['date'])           # assign the original dates as index
    )

date
2021-01-01   2022-02-01
2022-02-01   2022-03-01
2022-03-01   2022-04-01
2022-04-01   2022-05-01
2022-05-01          NaT
2022-06-01          NaT
Name: date, dtype: datetime64[ns]

Then, compute your operation (here GroupBy.max) and map the Series to your output:

data.groupby(['car_id', 'brand'])['date'].max().map(s)

output:

car_id  brand
1       a       2022-04-01
2       a              NaT
3       b              NaT
Name: date, dtype: datetime64[ns]
user_163434640
Related