基于Pandas实现N连续值与M连续值间的间隔距离计算

精简版需求 #

我有如下一列数据：

[2, 2, 0, 0, 0, 2, 2, 0, 3, 3, 3, 0, 0, 2, 2, 0, 0, 0, 0, 2, 2, 0, 0, 0, 3, 3, 3]
# 序列中可能存在更多数值，如4、5、6、7、8等

需要一个接收参数n、m的函数，当传入n=2、m=3时，能得到2连续值与3连续值之间的间隔距离，最终分组后的结果应为：

[6, 9]

详细需求 #

以下是测试用例代码：

dummy = [1,1,0,0,0,1,1,0,1,1,1,0,0,1,1,0,0,0,0,1,1,0,0,0,1,1,1]

df = pd.DataFrame({'a': dummy})

当前已编写的函数仅支持单个参数N（计算N连续值到另一N连续值的间隔），需要修改使其支持参数M，实现计算N连续值到M连续值的间隔。当前函数代码如下：

import pandas as pd
import numpy as np

def get_N_seq_stat(df, N=2, M=3):
    df["c1"] = (
        df.groupby(df.a.ne(df.a.shift()).cumsum())["a"]
        .transform("size")
        .where(df.a.eq(1), 0)
    )
    df["c2"] = np.where(df.c1.ne(N) , 1, 0)
    df["c3"] = df["c2"].ne(df["c2"].shift()).cumsum()

    result = df.loc[df["c2"] == 1].groupby("c3")["c2"].count().tolist()

    # if last N rows are not consequence shouldn't add last.
    if not (df["c1"].tail(N) == N).all():
        del result[-1]
    if not (df["c1"].head(N) == N).all():
        del result[0]
    return result    

预期与当前结果对比 #

• 预期结果：当设置N=2、M=3时，应返回[6,9]，对应序列（加粗部分为计算的间隔段）：

  dummy = [1,1,**0,0,0,1,1,0,**1,1,1,0,0,1,1,**0,0,0,0,1,1,0,0,0,**1,1,1]
  

• 当前结果：仅设置N=2时，函数返回[3, 6, 4]，对应序列：

  dummy = [1,1,**0,0,0,**1,1,**0,1,1,1,0,0,**1,1,**0,0,0,0,**1,1,0,0,0,1,1,1]
  

内容的提问来源于stack exchange，提问作者Yiffany