python - pandas 中的 groupby 具有必须保持状态的功能

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)

好处是，它运行完美! 缺点:在大型数据集(700 万个条目、25 万个不同的“a”值)上速度非常慢。

有什么更好的事情可以做吗？

Answer 1

您可以使用 groupby 和 concat 值查找 c 列的数字部分

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