使用pandas在同一索引的列中查找连续天的开始和结束日期

Question

我有一个数据框df：

df =

index  date        hats
A1     01-01-2020  5
A1     02-01-2020  10
A1     03-01-2020  16
A1     04-01-2020  16
A1     21-01-2020  9
A1     22-01-2020  8
A1     23-01-2020  7
A6     20-03-2020  5
A6     21-03-2020  5
A8     30-07-2020  12

在这里，前四行是连续的天数。我想知道数据框中所有此类连续天的开始日期和结束日期。如果在df 中的一系列明智的A8 索引中只有一天，那么开始日期和结束日期将相同。此外，我还想知道连续几天中df['hats'] 列中的最高值，并在单独的列high_hat 中返回其日期及其日期high_hat_date。如果连续几天有两个或更多相等的高值，则在新列num_hat中写入高值的出现次数，并在high_hat_date中写入第一次出现的日期。

上述数据框的示例输出如下：

index   start_date    end_date    high_hat    high_hat_date   num_hat
A1      01-01-2020    04-01-2020  16          03-01-2020      2
A1      21-01-2020    23-01-2020  9           21-01-2020      1
A6      20-03-2020    21-03-2020  5           20-03-2020      2
A8      30-07-2020    30-07-2020  12          30-07-2020      1     

非常感谢您在这方面的任何帮助。

Answer 1

首先使用pd.to_datetime 将date 列转换为pandas datetime 系列：

df['date'] = pd.to_datetime(df['date'], dayfirst=True)

然后使用：

g = df.groupby('index')['date'].diff().dt.days.ne(1).cumsum() # STEP A
m = df.groupby(['index', g])['hats'].transform('max').eq(df['hats']) # STEP B

df = df.assign(high_hats=df['hats'].mask(~m), high_date=df['date'].mask(~m)) # STEP C

dct = {'start_date': ('date', 'first'), 'end_date': ('date', 'last'), 'high_hat': ('hats', 'max'),
       'high_hat_date': ('high_date', 'first'), 'num_hats': ('high_hats', 'count')}
df1 = df.groupby(['index', g]).agg(**dct).reset_index().drop('date', 1) # STEP D

详情：

步骤 A：在 index 上使用 DataFrame.groupby 并在 date 上使用 groupby.diff 计算连续日期之间经过的天数，然后使用 Series.dt.days + Series.ne 和 Series.cumsum 创建分组系列g 将需要在连续日期对数据框进行分组。

# print(g)
0    1
1    1
2    1
3    1
4    2
5    2
6    2
7    3
8    3
9    4
Name: date, dtype: int64

步骤 B：在 index 和 g 上使用 DataFrame.groupby 并使用 groupby.transform 使用 max 转换列 hats，然后使用 Series.eq 将其与 hats 列等同以创建布尔值掩码m.

# print(m)
0    False
1    False
2     True
3     True
4     True
5    False
6    False
7     True
8     True
9     True
Name: hats, dtype: bool

步骤 C：接下来使用 DataFrame.assign 分配两个新列 high_hats 和 high_date，它们将在 STEP D 中用于计算 high_hat_date 和 num_hats。

# print(df)    
  index       date  hats  high_hats  high_date
0    A1 2020-01-01     5        NaN        NaT
1    A1 2020-01-02    10        NaN        NaT
2    A1 2020-01-03    16       16.0 2020-01-03
3    A1 2020-01-04    16       16.0 2020-01-04
4    A1 2020-01-21     9        9.0 2020-01-21
5    A1 2020-01-22     8        NaN        NaT
6    A1 2020-01-23     7        NaN        NaT
7    A6 2020-03-20     5        5.0 2020-03-20
8    A6 2020-03-21     5        5.0 2020-03-21
9    A8 2020-07-30    12       12.0 2020-07-30

步骤 D：在 index 和 g 上使用 DataFrame.groupby，并使用聚合字典 dct 聚合数据帧，该字典包含所有要应用的列及其对应的 agg 函数。

# print(df1)
  index start_date   end_date  high_hat high_hat_date  num_hats
0    A1 2020-01-01 2020-01-04        16    2020-01-03         2
1    A1 2020-01-21 2020-01-23         9    2020-01-21         1
2    A6 2020-03-20 2020-03-21         5    2020-03-20         2
3    A8 2020-07-30 2020-07-30        12    2020-07-30         1

Answer 2

这需要大量的数据转换：

1. 使用np.where()、.shift和.groupby + .transform创建一些中间列来计算组的数据范围。
2. 创建中间数据框df2 以计算更多指标，包括high_hat_date 和num_hat。这些计算需要关注最大值（即高帽），因此以这种方式创建这个新数据框更容易。
3. 将 df2 合并回 df1，只取所需的列并删除重复的行

代码：

import pandas as pd, numpy as np
df1=df.copy()
df1['date'] = pd.to_datetime(df1['date'], dayfirst=True)
df1['date_diff'] = df1['date'] - df1.shift()['date']
df1['date_first'] = ''
df1['date_first'] = np.where((df1['date_diff'].isnull()) |
                             ((df1['date_diff'] != '1 days') & (df1.shift()['date_diff'] == '1 days')),
                              'start_date', df1['date_first'])
df1['date_first'] = np.where((df1['date_diff'] == '1 days') & (df1.shift(-1)['date_diff'] != '1 days'),
                             'end_date', df1['date_first'])
df1['date_group'] = df1.groupby(df1['date_first'])['date_first'].transform('cumcount')
df1['date_group2'] = df1.groupby(df1['date_first'])['date_group'].transform('cumsum').replace(0,np.nan).ffill().astype(int)
df1['start_date'] = df1.groupby('date_group2')['date'].transform('min')
df1['end_date'] = df1.groupby('date_group2')['date'].transform('max')
df1['high_hat'] = df1.groupby(df1['date_group2'])['hats'].transform('max')
df2 = df1.loc[df1['high_hat'] == df1['hats']]
df2['high_hat_date'] = df2.groupby('date_group2')['date'].transform('first')
df2['num_hat'] = df2.groupby('date_group2')['hats'].transform('count')
df2 = df2.drop_duplicates(subset='date_group2')
df1 = pd.merge(df1, df2[['date_group2', 'high_hat_date', 'num_hat']], how='outer', on=['date_group2'])
df1 = df1[['index', 'start_date', 'end_date', 'high_hat', 'high_hat_date', 'num_hat']].drop_duplicates()
df1

输出：

    index   start_date  end_date    high_hat    high_hat_date   num_hat
0   A1      2020-01-01  2020-01-04  16          2020-01-03      2
4   A1      2020-01-21  2020-01-23  9           2020-01-21      1
7   A6      2020-03-20  2020-03-21  5           2020-03-20      2
9   A8      2020-07-30  2020-07-30  12         2020-07-30       1