Pandas DataFrame：查找兩列相等/相同的特定長度序列的索引值

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是相等A ，即我想提取索引值必須是連續的（長度> = n的切片），其中A和B相等。

所以在這種情況下，對於n = 2，它應該是“twos”和“fours”的索引：

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個連續的相等元素？

我想必須有一些我目前看不到的奇特的組合方式..

Answer 1

根據您的描述，我不清楚為什么1和5將被排除在您的結果之外，因為它們每個都包含2個或更多連續的行，並且具有A和B的匹配值。

但是，下面的解決方案仍然有用，我相信您可以根據自己的需要進行修改。 它首先過濾數據幀以匹配A列和B列中A值（ 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

上面的數據框是為了確保它符合您的期望。 然后只提取索引：

>>> 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 ，則會得到預期結果。