сгруппировать в пандах с функцией, которая должна сохранять состояние

Question

У меня есть следующий фрейм данных

df = pd.DataFrame({'a': ['A', 'A', 'A', 'B', 'B', 'B', 'B'],
                   'b': [  1,   2,   4,   1,   2,   3,   4]})

Мне нужна функция, которая выдает следующее определение фрейма данных:

df = pd.DataFrame({'a': [  'A',   'A',   'A',   'B',   'B',   'B',   'B'],
                   'b': [    1,     2,     4,     1,     2,     3,     4],
                   'c': ['A_0', 'A_0', 'A_1', 'B_0', 'B_0', 'B_0', 'B_0']})

Логика такова, что для каждого значения 'a'(каждая группа), я создаю значение «c», которое можно описать как «непрерывный» ряд значений «b».

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

def detection(dataset):
    def detect(series, avalue):
        _id = 0
        start = True
        visits = []
        prev_ = None
        for h in series:
            if start:
                start = False
                prev_ = h
            else:
                if h - prev_ > 1:
                    _id += 1
                prev_ = h
            visits.append(f"{avalue}_{_id}")
        return visits

    res = []
    gb = dataset.groupby("a")
    for avalue in gb.groups:
        dd = gb.get_group(avalue)
        dd["VISIT_ID"] = detect(dd["b"], avalue)
        res.append(dd)
    return pd.concat(res, axis=0)

Хорошо, это работает отлично!Плохо: он очень медленный для большого набора данных (7 миллионов записей, 250 тыс. Различных значений 'a'.

Есть ли что-нибудь лучше?

Answer 1

Вы можете найти числовую часть столбца c, используя значения groupby и concat

df['c'] = df.groupby('a').b.apply(lambda x: (x.diff() > 1).cumsum())
df['c'] = df['a'] + '_' + df['c'].astype(str)


    a   b   c
0   A   1   A_0
1   A   2   A_0
2   A   4   A_1
3   B   1   B_0
4   B   2   B_0
5   B   3   B_0
6   B   4   B_0