Pandas分组计算交通数据时仅第一组vehicle_real求和正确问题排查

问题场景 #

处理交通数据时，按time_interval_code分组计算vehicle_real字段求和，仅第一组(1,2,3,4)结果正确，其余组数值异常，重复的时段未被纳入新组求和范围。

背景说明：上午共划分8个15分钟时段，time_interval_code用1-8对应（07:00-07:15为1，依此类推至8）。需统计每个路口(junction_id)、来车方向(source_direction)的小时最大流量——即每次对连续4个时段的vehicle_real求和，共需计算5组连续时段：(1,2,3,4)、(2,3,4,5)、(3,4,5,6)、(4,5,6,7)、(5,6,7,8)。

原代码如下：

import pandas as pd
data = pd.read_excel("traffic.xlsx")

# Create a DataFrame from the list of data
df = pd.DataFrame(data)

# Define a function to get the morning groups for each time interval code
def get_morning_group(time_interval_code):
  morning_groups = [(1, 2, 3, 4), (2, 3, 4, 5), (3, 4, 5, 6), (4, 5, 6, 7), (5, 6, 7, 8)]

  for group in morning_groups:
    if time_interval_code in group:
      return group

# Add a new column to the DataFrame that contains the morning groups for each time interval code
df['morning_groups'] = df['time_interval_code'].apply(get_morning_group)

# Group data by values
grouped_data = df.groupby(['junction_id', 'source_direction', 'morning_groups'])

# Calculate the sum of the vehicles_real values for each group
grouped_data = grouped_data['vehicles_real'].sum()

# Convert the grouped data back into a DataFrame
df = grouped_data.reset_index()

# Create the pivot table
pivot_table = df.pivot_table(index=['junction_id', 'source_direction'], columns=['morning_groups'], values='vehicles_real')

# Save the pivot table to a new Excel file
pivot_table.to_excel('max_flow_rate.xlsx')

数据规模：约14万条记录，每个路口+来车方向组合在所有时段均有vehicle_real数据。

问题原因 #

原代码的get_morning_group函数逻辑错误：当一个time_interval_code属于多个组时（比如2同时属于(1,2,3,4)和(2,3,4,5)），函数仅返回第一个匹配的组，导致后续组无法获取该时段的数据。例如时段2只会被分到第一组，不会出现在第二组中，最终第二组求和时缺少时段2的数据，结果自然异常。

解决方案 #

需要为每个时段生成所有所属的组，而非仅返回第一个匹配组，通过explode展开分组实现：

修正后的代码 #

import pandas as pd

# 读取数据
df = pd.read_excel("traffic.xlsx")

# 定义所有需要计算的连续时段组
morning_groups = [(1,2,3,4), (2,3,4,5), (3,4,5,6), (4,5,6,7), (5,6,7,8)]

# 构建时段与组的映射表：每个时段对应所有包含它的组
group_mapping = {}
for group in morning_groups:
    for code in group:
        if code not in group_mapping:
            group_mapping[code] = []
        group_mapping[code].append(group)

# 将每个时段对应的所有组添加到DataFrame，并展开成行
df['morning_groups'] = df['time_interval_code'].map(group_mapping)
df = df.explode('morning_groups')

# 分组求和
grouped_data = df.groupby(['junction_id', 'source_direction', 'morning_groups'])['vehicles_real'].sum().reset_index()

# 生成透视表并保存
pivot_table = grouped_data.pivot_table(
    index=['junction_id', 'source_direction'],
    columns=['morning_groups'],
    values='vehicles_real'
)
pivot_table.to_excel('max_flow_rate.xlsx')

关键改进点 #

• 重构group_mapping：确保每个时段对应所有包含它的组（比如时段2对应[(1,2,3,4), (2,3,4,5)]）
• 使用explode('morning_groups')：将每个时段的多个组拆分成单独行，保证分组时每个组都能包含所有属于它的时段数据
• 后续分组求和、透视表逻辑保持不变，但数据基础已正确覆盖所有组的时段

额外优化（可选） #

若需直接计算小时最大流量，可在透视表后添加最大值列：

# 添加最大流量列
pivot_table['max_hour_flow'] = pivot_table.max(axis=1)
pivot_table.to_excel('max_flow_rate_with_max.xlsx')

内容的提问来源于stack exchange，提问作者mixx