pandas中groupby和shift结合实现相邻行的计算

⭐️ pandas中的groupby和shift

在 Pandas 中，groupby 操作通常用于按某个列或多个列分组，然后对每个组应用某些聚合、变换或过滤操作。shift 方法用于对数据进行移位操作，它将数据沿某个轴（通常是行）进行平移。

将 groupby 和 shift 结合使用，可以对分组后的数据进行移位操作。这在需要计算组内变化或生成滞后变量（如时间序列数据中的滞后值）时非常有用。

这里给出一个示例，实现相邻两行中时间的计算。

⭐️ 准备数据

这里准备两个方法，一个方法随机生成一些时间，

def generate_random_times(start_time, end_time, num_times):
    """
    生成指定数量的随机时间，精确到秒。
    
    参数：
    start_time: datetime, 起始时间
    end_time: datetime, 结束时间
    num_times: int, 生成时间的数量
    
    返回：
    list of datetime, 随机生成的时间列表
    """
    delta = end_time - start_time
    random_times = []
    
    for _ in range(num_times):
        random_second = random.randint(0, int(delta.total_seconds()))
        random_time = start_time + timedelta(seconds=random_second)
        random_times.append(random_time)
    
    return random_times

另一个方法随机选取每行数据的组别

def generate_random_groups(num_groups=N):
    group_list = ["A", "B", "C"]
    # 随机选择多个值，允许重复
    _groups = [random.choice(group_list) for _ in range(num_groups)]
    return _groups

⭐️ 完整代码

完整代码如下

import random
import pandas as pd
from datetime import datetime, timedelta

# 数据量
N = 20

def generate_random_times(start_time, end_time, num_times):
    """
    生成指定数量的随机时间，精确到秒。
    
    参数：
    start_time: datetime, 起始时间
    end_time: datetime, 结束时间
    num_times: int, 生成时间的数量
    
    返回：
    list of datetime, 随机生成的时间列表
    """
    delta = end_time - start_time
    random_times = []
    
    for _ in range(num_times):
        random_second = random.randint(0, int(delta.total_seconds()))
        random_time = start_time + timedelta(seconds=random_second)
        random_times.append(random_time)
    
    return random_times


def generate_random_groups(num_groups=N):
    group_list = ["A", "B", "C"]
    # 随机选择多个值，允许重复
    _groups = [random.choice(group_list) for _ in range(num_groups)]
    return _groups

# 时间
start_time = datetime(2024, 1, 1, 0, 0, 0)
end_time = datetime(2024, 12, 31, 23, 59, 59)
random_times = generate_random_times(start_time, end_time, N)
# 组别
random_groups = generate_random_groups()

# 创建示例数据
data = {
    'group': random_groups,
    'date': random_times
}
df = pd.DataFrame(data)
# 按组别和时间排序，并重置索引
sorted_df = df.sort_values(by=['group', 'date']).reset_index(drop=True)
# 按 'group' 列进行分组，然后对每个组的时间进行移位
sorted_df['shifted_date'] = sorted_df.groupby('group')['date'].shift(1)
# 删除缺失值，并重置索引
result_df = sorted_df.dropna().reindex()
# 计算相连两行的时间间隔
result_df['lag'] = pd.to_datetime(result_df['date']) - pd.to_datetime(result_df['shifted_date'])
result_df['lag_seconds'] = result_df['lag'].dt.total_seconds()
result_df = result_df.reset_index(drop=True)
# 输出结果
print(result_df)

运行结果如下，由于时间和组别都存在随机因素，所以每次执行结果会不同

笔者水平有限，若有不对的地方欢迎评论指正！