如何高效计算洛杉矶地铁站点空闲时间？现有代码存场景遗漏

问题 #

我们在大学项目中处理洛杉矶地铁160万行共享单车交易数据，字段包含transaction_id、start_time、start_station_id、end_time、end_station_id、bike_id。需求是计算各站点的自行车空闲时间——即同一站点内，某辆自行车被归还后到下一次被租用的时间间隔。

现有基于Pandas的代码无法处理重叠时段的场景，以下是简化测试数据：

transaction_id,start_time, start_station_id, end_time, end_station_id
0,2017-01-01 10:00:00, 3027, 2017-01-01 10:10:00, 3027
1,2017-01-01 10:01:00, 3027, 2017-01-01 10:05:00, 3027
2,2017-01-01 10:03:00, 3027, 2017-01-01 10:20:00, 3027
3,2017-01-01 10:07:00, 3027, 2017-01-01 10:09:00, 3027
4,2017-01-01 10:12:00, 3027, 2017-01-01 10:23:00, 3027

预期输出需新增idle_time列，无空闲时间则为none，否则为时间间隔：

transaction_id,start_time, start_station_id, end_time, end_station_id,idle_time
0,2017-01-01 10:00:00, 3027, 2017-01-01 10:10:00, 3027,none
1,2017-01-01 10:01:00, 3027, 2017-01-01 10:05:00, 3027,none
2,2017-01-01 10:03:00, 3027, 2017-01-01 10:20:00, 3027,none
3,2017-01-01 10:07:00, 3027, 2017-01-01 10:09:00, 3027,0 days 00:02:00
4,2017-01-01 10:12:00, 3027, 2017-01-01 10:23:00, 3027,0 days 00:02:00

当前使用的代码如下：

import pandas as pd
from datetime import datetime, timedelta
csv_file = '../Data_Cleanup/outCSV/Clean_Metro_Set.csv'
metro = pd.read_csv(csv_file)
metro['start_time'] = pd.to_datetime(metro['start_time'])

metro ['end_time'] = pd.to_datetime(metro['end_time'])

metro = metro.sort_values(by='start_time')
metro['idle_time'] = None

BigDict = {}  # Initialize an empty dictionary to store station IDs and end times

# Iterate through each row in the DataFrame
for i, row in metro.iterrows():
    current_start_station_id = metro.loc[i, "start_station_id"]
    current_end_station_id = metro.loc[i, "end_station_id"]
    current_start_time = metro.loc[i, "start_time"]
    current_end_time = metro.loc[i, "end_time"]

    try:
        # Check if current start time is greater than the stored end time for the start station ID
        if current_start_time > BigDict[str(current_start_station_id)]:
            # Calculate idle time as the difference between the current start time and the stored end time
            metro.at[i, "idle_time"] = current_start_time - BigDict[str(current_start_station_id)]
    except KeyError:
        pass

    # Store the current end time for the end station ID in BigDict
    try:
        if current_end_time > BigDict[str(current_end_station_id)]:
            BigDict[str(current_end_station_id)] = current_end_time
    except KeyError:
        BigDict[str(current_end_station_id)] = current_end_time

需要在保证计算效率的前提下，修正代码以正确计算空闲时间。

解决方案 #

问题根源 #

原代码的核心错误是：用字典存储站点的最晚还车时间，但实际场景中站点可能同时有多辆自行车，新租车请求应匹配的是该站点中早于当前租车时间的最近一次还车记录。原代码仅保留最晚还车时间，会忽略更早的可用还车记录，导致重叠时段的空闲时间计算错误。同时，iterrows逐行迭代的方式处理百万级数据效率极低。

优化后的代码 #

import pandas as pd

# 读取并预处理数据
csv_file = '../Data_Cleanup/outCSV/Clean_Metro_Set.csv'
metro = pd.read_csv(csv_file)
metro['start_time'] = pd.to_datetime(metro['start_time'])
metro['end_time'] = pd.to_datetime(metro['end_time'])

# 提取所有还车记录，按站点和还车时间排序
returns = metro[['end_station_id', 'end_time']].rename(
    columns={'end_station_id': 'station_id', 'end_time': 'return_time'}
).sort_values(by=['station_id', 'return_time'])

# 原始数据按站点+租车时间排序，满足merge_asof的匹配要求
metro_sorted = metro.sort_values(by=['start_station_id', 'start_time']).reset_index(drop=True)

# 使用merge_asof高效匹配最近的、早于当前租车时间的还车记录
merged = pd.merge_asof(
    metro_sorted,
    returns,
    left_on='start_time',
    right_on='return_time',
    left_by='start_station_id',
    right_by='station_id',
    direction='backward'  # 匹配小于等于当前租车时间的最近还车记录
)

# 计算空闲时间，无匹配则设为'none'
metro_sorted['idle_time'] = merged['start_time'] - merged['return_time']
metro_sorted['idle_time'] = metro_sorted['idle_time'].apply(
    lambda x: str(x) if pd.notnull(x) else 'none'
)

# 恢复原交易ID排序（可选，根据需求调整）
metro_sorted = metro_sorted.sort_values(by='transaction_id').reset_index(drop=True)

print(metro_sorted)

代码说明 #

1. 效率保障：使用Pandas的merge_asof函数，该函数基于排序后的二分查找实现匹配，时间复杂度为O(n log n)，处理百万级数据的速度远快于逐行迭代。
2. 逻辑正确性：
   • 单独提取还车记录并排序，确保每个站点的还车时间按时间顺序排列。
   • direction='backward'参数保证匹配到当前租车时间之前的最近一次还车记录，完全符合空闲时间的定义。
3. 灵活性：最后可选择恢复原交易ID排序，也可保留按站点+时间排序的结果，适配不同输出需求。

验证结果 #

用示例输入测试，输出与预期完全一致：

• 交易0、1、2无早于其租车时间的还车记录，idle_time为none；
• 交易3匹配到交易1的10:05还车记录，空闲时间为2分钟；
• 交易4匹配到交易0的10:10还车记录，空闲时间为2分钟。

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