pandas 列表按间隔分组

看过来

《pandas 教程》持续更新中，提供建议、纠错、催更等加作者微信: gairuo123（备注：pandas教程）和关注公众号「盖若」ID: gairuo。跟作者学习，请进入 Python学习课程。欢迎关注作者出版的书籍：《深入浅出Pandas》和《Python之光》。

有一个列表[3,6,22,0,25,1,7,5,2,13,14,26,27]，从小到大排序后，将数字连续间隔小于为1或相等的数放在一起，形成一个新的列表，即[3,6,22,0,25,1,7,5,2,13,14,26,27]。

解决该问题，我们有两个思路，一个是排序后将间隔大于1的节点找出来，然后用这些节点进行分箱操作，分箱后再进行分组，另外一个是找到上述节点后，将非节点当缺失值处理，用fillna全填充为节点的值作为分组标识，再按这些标识进行分组。接下来我们先实现第一个思路，先将数据放入DataFrame，然后排好序：

import pandas as pd
import numpy as np

numbers = [3,6,22,0,25,1,7,5,2,13,14,26,27] 
df = pd.DataFrame({'n': numbers})
df = df.sort_values('n', ignore_index=True)
df.head()
'''
     n
0    0
1    1
2    2
3    3
4    5
5    6
6    7
7   13
8   14
9   22
10  25
11  26
12  27
'''

增加一个与前一个数据的差值列：

df.assign(diff=df.n.diff(1))
'''
     n  diff
0    0   NaN
1    1   1.0
2    2   1.0
3    3   1.0
4    5   2.0
5    6   1.0
6    7   1.0
7   13   6.0
8   14   1.0
9   22   8.0
10  25   3.0
11  26   1.0
12  27   1.0
'''

将大于1的筛选出来，n列的数字就是这些分隔的节点：

df.assign(diff=df.n.diff(1)).query('diff>1')
'''
     n  diff
4    5   2.0
7   13   6.0
9   22   8.0
10  25   3.0
'''

将这个节点输出为列表，赋值给变量bins，用于后续的分箱中：

bins = (
    df.assign(diff=df.n.diff(1))
    .query('diff>1')
    .n
    .to_list()
)
bins
# [5, 13, 22, 25]

开始操作分箱，分箱的目标为n列，箱子不包含右边的值，同时补充0：

df.set_axis(pd.cut(df.n, [0]+bins, right=False))
'''
               n
n
[0.0, 5.0)     0
[0.0, 5.0)     1
[0.0, 5.0)     2
[0.0, 5.0)     3
[5.0, 13.0)    5
[5.0, 13.0)    6
[5.0, 13.0)    7
[13.0, 22.0)  13
[13.0, 22.0)  14
[22.0, 25.0)  22
NaN           25
NaN           26
NaN           27
'''

可以看到，索引处已是对数字分组的箱子标识，再对数据进行分组，分组后再输出最终列表：

(
    df.set_axis(pd.cut(df.n, [0]+bins, right=False))
    .groupby(lambda x:x, dropna=False)
    .apply(lambda x: x.n.to_list())
    .to_list()
)
# [[0, 1, 2, 3], [5, 6, 7], [13, 14], [22], [25, 26, 27]]

再看看第二个思路，先增加一个辅助列，用于标识的初始值：

(
    pd.DataFrame({'n': [3,6,22,0,25,1,7,5,2,13,14,26,27] })
    .sort_values('n', ignore_index=True)
    .assign(sign=lambda x: x.index)
)
'''
     n  sign
0    0     0
1    1     1
2    2     2
3    3     3
4    5     4
5    6     5
6    7     6
7   13     7
8   14     8
9   22     9
10  25    10
11  26    11
12  27    12
'''

在sign列上计算和前一个值的差值是否大于1，如果相差超过1取原值，否则设置为缺失值，以方便接下来的缺失值填充：

(
    pd.DataFrame({'n': [3,6,22,0,25,1,7,5,2,13,14,26,27] })
    .sort_values('n', ignore_index=True)
    .assign(sign=lambda x: x.index)
    .assign(sign=lambda x: np.where(x.n.diff(1)>1, x.sign, np.NaN))
)
'''
     n  sign
0    0   NaN
1    1   NaN
2    2   NaN
3    3   NaN
4    5   4.0
5    6   NaN
6    7   NaN
7   13   7.0
8   14   NaN
9   22   9.0
10  25  10.0
11  26   NaN
12  27   NaN
'''

操作缺失值填充，填充前一个有效值，同时将NaN填充为0：

(
    pd.DataFrame({'n': [3,6,22,0,25,1,7,5,2,13,14,26,27] })
    .sort_values('n', ignore_index=True)
    .assign(sign=lambda x: x.index)
    .assign(sign=lambda x: np.where(x.n.diff(1)>1, x.sign, np.NaN))
    .fillna(method='ffill')
    .fillna(0)
)
'''
     n  sign
0    0   0.0
1    1   0.0
2    2   0.0
3    3   0.0
4    5   4.0
5    6   4.0
6    7   4.0
7   13   7.0
8   14   7.0
9   22   9.0
10  25  10.0
11  26  10.0
12  27  10.0
'''

以上就完成了分组标识的建立，最后再完成按标识分组和转列表操作：

(
    pd.DataFrame({'n': [3,6,22,0,25,1,7,5,2,13,14,26,27] })
    .sort_values('n', ignore_index=True)
    .assign(sign=lambda x: x.index)
    .assign(sign=lambda x: np.where(x.n.diff(1)>1, x.sign, np.NaN))
    .fillna(method='ffill')
    .fillna(0)
    .groupby('sign')
    .apply(lambda x: x.n.to_list())
    .to_list()
)
# [[0, 1, 2, 3], [5, 6, 7], [13, 14], [22], [25, 26, 27]]

这样就完成了这个需求，看似代码略多些，但整体逻辑比较清晰。

pandas 列表按间隔分组

相关内容