首页 > 解决方案 > 根据开始和停止列对数据框进行分组

问题描述

我想根据开始和停止列剪切/分组熊猫数据框,但仅限于开始->停止的情况。

我想要从“开始”非零值到“停止”非零值的索引范围。但仅当“开始”非零值后跟“停止”非零值时。从上到下遍历索引

我附上了一些代码,创建了问题的简化版本和相应的图像。

col1 = np.zeros(10)
col2 = np.zeros(10)
col1[[0, 1, 5, 8]] = 1
col2[[3, 6, 7, 9]] = 1

df = pd.DataFrame({'start': col1, 'stop': col2})

示例数据框

所需的输出将索引分组,有点像:[(1,2,3),(5,6),(8,9)]

以防万一这会简化事情的附加信息:

  1. 合并列会很好。
  2. 我的原始数据框有一个 pd.TimedeltaIndex。

所需结果的视觉澄清: 视觉澄清

标签: pythonpandasdataframe

解决方案


首先,我们需要查看 and 的区间startstop找出哪些是“有效”区间结束:

>>> ends = df.index.to_series().where(df['stop'].ne(0))
>>> starts = df.index.to_series().where(df['start'].ne(0))
>>> ends
0    NaN
1    NaN
2    NaN
3    3.0
4    NaN
5    NaN
6    6.0
7    7.0
8    NaN
9    9.0
dtype: float64
>>> starts
0    0.0
1    1.0
2    NaN
3    NaN
4    NaN
5    5.0
6    NaN
7    NaN
8    8.0
9    NaN
dtype: float64

现在我们可以尝试为每个有效开始获取下一个有效结束:

>>> next_end = ends.bfill().rename('end')
>>> valid_starts = starts.dropna().rename('start')
>>> candidates = valid_starts.to_frame().join(next_end, how='left')
>>> candidates
   start  end
0    0.0  3.0
1    1.0  3.0
5    5.0  6.0
8    8.0  9.0

在这里,我们看到从 0 开始的间隔存在问题:另一个间隔稍后开始(在 1),因此 [0, 3] 无效,我们应该只保留 [1, 3]。这可以通过 groupby + max 来完成,例如:

>>> intervals = candidates.groupby('end')['start'].max().reset_index().astype(int)
>>> intervals
   end  start
0    3      1
1    6      5
2    9      8

最后从端点生成索引列表很容易:

>>> intervals.agg(lambda s: list(range(s['start'], s['end'] + 1)), axis='columns')
0    [1, 2, 3]
1       [5, 6]
2       [8, 9]
dtype: object

推荐阅读