Как группировать по одному столбцу и по разнице соседних значений == 1 в другом столбце?

Имеется таблица:

data1 = pd.DataFrame({'store_id': [4168621, 4168621, 4168621, 4168621, 4168621, 4168624, 4168624, 4168626, 4168626],
                          'period_id': [191, 192, 194, 195, 196, 191, 192, 193, 194] })

data1:

   store_id  period_id
0   4168621        191
1   4168621        192
2   4168621        194
3   4168621        195
4   4168621        196
5   4168624        191
6   4168624        192
7   4168626        193
8   4168626        194

Необходимый выход:

data2 = pd.DataFrame({'store_id': [4168621, 4168621, 4168621, 4168621, 4168621, 4168624, 4168624, 4168626, 4168626],
                          'period_id': [191, 192, 194, 195, 196, 191, 192, 193, 194],
                          'promo_id': [1, 1, 2, 2, 2, 3, 3, 4, 4],
                          'promo_cw': [0, 1, 0, 1, 2, 0, 1, 0, 1] })

data2:

   store_id  period_id  promo_id  promo_cw
0   4168621        191         1         0
1   4168621        192         1         1
2   4168621        194         2         0
3   4168621        195         2         1
4   4168621        196         2         2
5   4168624        191         3         0
6   4168624        192         3         1
7   4168626        193         4         0
8   4168626        194         4         1

Рабочий код для строк под индексами 0 - 6:

data1['promo_id'] = ((data1.period_id - data1.period_id.shift()).fillna(1) != 1).cumsum() + 1
data1['promo_cw'] = data1.groupby('promo_id').cumcount() 
   store_id  period_id  promo_id
0   4168621        191         1
1   4168621        192         1
2   4168621        194         2
3   4168621        195         2
4   4168621        196         2
5   4168624        191         3
6   4168624        192         3
7   4168626        193         3
8   4168626        194         3

Рабочий код для строк под индексами 7 - 8:

data1['promo_id'] = ((data1.store_id == data1.store_id.shift()).fillna(1) == 0).cumsum() + 1
data1['promo_cw'] = data1.groupby('promo_id').cumcount() 
   store_id  period_id  promo_id
0   4168621        191         2
1   4168621        192         2
2   4168621        194         2
3   4168621        195         2
4   4168621        196         2
5   4168624        191         3
6   4168624        192         3
7   4168626        193         4
8   4168626        194         4

Через символ "|" объединить не получилось.

Как можно объединить эти фрагменты кода, чтобы он работал для всей таблицы?


Ответы (1 шт):

Автор решения: MaxU

Попробуйте так:

data1["promo_id"] = (
    (data1["store_id"] != data1["store_id"].shift().bfill())
    |
    (data1["period_id"].diff().fillna(1).ne(1))
).cumsum().add(1)
data1['promo_cw'] = data1.groupby("promo_id").cumcount()

результат:

In [77]: data1
Out[77]:
   store_id  period_id  promo_id  promo_cw
0   4168621        191         1         0
1   4168621        192         1         1
2   4168621        194         2         0
3   4168621        195         2         1
4   4168621        196         2         2
5   4168624        191         3         0
6   4168624        192         3         1
7   4168626        193         4         0
8   4168626        194         4         1

PS если хотите разобраться как это работает, вам придется запускать отдельные кусочки и смотреть, что из этого получается, а потом последовательно комбинировать их и опять смотреть на промежуточные результаты...

→ Ссылка