Как упорядочить строку на основе другой строки?

Question

Я пытаюсь преобразовать формулу из Excel в pandas.

. DataFrame выглядит следующим образом:

Column A    Column B 
H  
H  
H  
J  
J  
J  
J  
K  
K  

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

Column A     Column B
H            1
H            2
H            3
J            1
J            2
J            3
J            4
K            1
K            2

В Excel формула будет =IF(A2<>A1,1,B1+1)

Как я могу применить эту формулу в pandas?

Answer 1

Это можно сделать, используя следующий векторизованный метод:

Код:

>>> df = pd.DataFrame({'A':['H', 'H', 'H', 'J', 'J', 'J', 'J', 'K', 'K']})
>>> df['B'] = df.groupby((df['A'].shift(1) != df['A']).cumsum()).cumcount() + 1

Выход:

>>> df
   A  B
0  H  1
1  H  2
2  H  3
3  J  1
4  J  2
5  J  3
6  J  4
7  K  1
8  K  2

Объяснение:

Сначала мы используем df['A'].shift(1) != df['A'] для сравнения столбца A со столбцом A, смещенным на 1. Это дает:

>>> df['A'] != df['A'].shift(1)
0     True
1    False
2    False
3     True
4    False
5    False
6    False
7     True
8    False
Name: A, dtype: bool

Далее мы используем cumsum(), чтобы вернуть накопленную сумму по этому столбцу. Это дает нам:

>>> (df['A'] != df['A'].shift(1)).cumsum()
0    1
1    1
2    1
3    2
4    2
5    2
6    2
7    3
8    3
Name: A, dtype: int32

Теперь мы можем использовать GroupBy.cumcount() как обычно, чтобы перечислять каждый элемент в порядке возрастания, добавляя 1, чтобы начать индекс с 1. Обратите внимание, что мы можем просто используйте

df.groupby('A').cumcount()

Потому что, если, например, у нас было:

>>> df
   A
0  H
1  H
2  H
3  J
4  J
5  J
6  J
7  K
8  K
9  H

Это дало бы нам:

>>> df.groupby('A').cumcount() + 1
0    1
1    2
2    3
3    1
4    2
5    3
6    4
7    1
8    2
9    4
dtype: int64

Обратите внимание, что последняя строка 4, а не 1, как ожидалось.