python - Pandas 数据帧 : Find index values for sequences of a certain length where two columns are equal/identical

Question

我有一个 pandas DataFrame 定义如下:

# -*- coding: utf-8 -*-
import datetime as dt
import pandas as pd


data = [[1, 1], [1, 1], [2, 2], [2, 2], [2, 2], [3, 3], [4, 4], [4, 4],
        [4, 4], [5, 5], [5, 5]]
df = pd.DataFrame(data, columns=['A', 'B'])
df.index = pd.date_range(dt.datetime(2012, 1, 1), periods=len(df), freq='H')

print(df)

和产量:

                 A  B
2012-01-01 00:00:00  1  1
2012-01-01 01:00:00  1  1
2012-01-01 02:00:00  2  2
2012-01-01 03:00:00  2  2
2012-01-01 04:00:00  2  2
2012-01-01 05:00:00  3  3
2012-01-01 06:00:00  4  4
2012-01-01 07:00:00  4  4
2012-01-01 08:00:00  4  4
2012-01-01 09:00:00  5  5
2012-01-01 10:00:00  5  5

现在我正在尝试获取 A 列和 B 列相等并且至少(或恰好也足够)n 个连续行(此处为小时)在 A< 中相等的行的索引 和 B 即我想提取必须连续的索引值(长度 >= n 的切片)，其中 A 和 B 相等。

所以在这种情况下，对于 n=2，它应该是“双”和“四”的索引:

2012-01-01 02:00:00
2012-01-01 03:00:00
2012-01-01 04:00:00
2012-01-01 06:00:00
2012-01-01 07:00:00
2012-01-01 08:00:00

只获取 A 和 B 相等的行的索引很简单。

但是我怎样才能只获取 n 个连续的相等元素呢？

我想一定有一些我目前没有看到的花哨的 groupby 方法..

Answer 1

根据您的描述，我不清楚为什么 1 和 5 会被排除在您的结果之外，因为它们每个都包含 2 个或更多连续行，并且 A 和 B 的值相匹配。

但是，下面的解决方案应该仍然有用，而且我相信您可以修改它以满足您的需要。它首先过滤数据帧以匹配列 A 和 B (df_matching) 中的值。然后，它使用 shift-cumsum 模式对连续匹配值进行分组，然后对 n 进行过滤。

n = 2
df_matching = df[df.A == df.B]
gb = df_matching.groupby((df_matching.A != df_matching.A.shift()).cumsum())
df_target = gb.filter(lambda x: len(x) >= n)

>>> df_target
                     A  B
2012-01-01 00:00:00  1  1
2012-01-01 01:00:00  1  1
2012-01-01 02:00:00  2  2
2012-01-01 03:00:00  2  2
2012-01-01 04:00:00  2  2
2012-01-01 06:00:00  4  4
2012-01-01 07:00:00  4  4
2012-01-01 08:00:00  4  4
2012-01-01 09:00:00  5  5
2012-01-01 10:00:00  5  5

上面的dataframe是为了保证符合你的预期。然后只提取索引:

>>> df_target.index
DatetimeIndex(['2012-01-01 00:00:00', '2012-01-01 01:00:00',
               '2012-01-01 02:00:00', '2012-01-01 03:00:00',
               '2012-01-01 04:00:00', '2012-01-01 06:00:00',
               '2012-01-01 07:00:00', '2012-01-01 08:00:00',
               '2012-01-01 09:00:00', '2012-01-01 10:00:00'],
              dtype='datetime64[ns]', freq=None)

请注意，如果 n=3，您会得到预期的结果。