根据 id 的给定值在 Python 中创建排名

Question

I have this dataset:

dic = {'id':[1,1,1,1,1,2,2,2,2], 'sales': [100.00, 200.00, 300.00, 400.00, 500.00, 100.00, 200.00, 300.00, 400.00], 'year_month': [202201, 202202, 0, 202204, 202205, 202201, 202202, 202203, 0]}
df = pd.DataFrame(dic)

Output:

   id   sales   year_month
0   1   100.0   202201
1   1   200.0   202202
2   1   300.0   0
3   1   400.0   202204
4   1   500.0   202205
5   2   100.0   202201
6   2   200.0   202202
7   2   300.0   202203
8   2   400.0   0

I want to increases 1 after year_month zero and decreases 1 before zero, per ID, like that:

    id  sales   year_month rank
0   1   100.0   202201     -2
1   1   200.0   202202     -1
2   1   300.0   0           0
3   1   400.0   202204      1
4   1   500.0   202205      2
5   2   100.0   202201     -3
6   2   200.0   202202     -2
7   2   300.0   202203     -1
8   2   400.0   0           0

How do I Create the rank column?

Answer 1

I came up with this. It seems more complicated than it's supposed to do, but it still works

difference = [(df[(df.id == id) & (df.year_month == year)].index - df[(df.id == id) & (df.year_month == 0)].index)[0] for id in df.id.unique() for year in df[df.id == id].year_month]

df['new'] = difference

which indeed gives

0   1   100.0   202201  -2
1   1   200.0   202202  -1
2   1   300.0   0        0
3   1   400.0   202204   1
4   1   500.0   202205   2
5   2   100.0   202201  -3
6   2   200.0   202202  -2
7   2   300.0   202203  -1
8   2   400.0   0        0

Answer 2

Given default index, sorted vals in id , and sorted vals in year_month ( 0 replacing a sorted val & always min for each group), you can simply do:

df['rank'] = df.index - df.groupby('id')['year_month'].transform('idxmin')
print(df)

   id  sales  year_month  rank
0   1  100.0      202201    -2
1   1  200.0      202202    -1
2   1  300.0           0     0
3   1  400.0      202204     1
4   1  500.0      202205     2
5   2  100.0      202201    -3
6   2  200.0      202202    -2
7   2  300.0      202203    -1
8   2  400.0           0     0