如何基于5分钟时间差与前置条件为DataFrame标记visit类型

问题背景 #

原有规则仅判断当前行的上一行是否满足：同公司（company）且类型（type）为home，若是则将当前行的source标记为visit。但该规则存在缺陷：当中间插入其他公司的记录时，无法识别到同公司更早的、时间差在5分钟内的home行（比如示例中行2属于apple，最近的apple home行是行0，时间差4分钟，但因中间隔了tesla的行，原规则无法标记）。

需要修改规则为：只要当前行与同公司最近的type为home的记录的时间差≤5分钟，就将当前行的source标记为visit。

示例输入 #

source datetime location  type  start  company 
0          10:00    london    home  1       apple
1          10:03    unknown                 tesla
2          10:04    France                  apple
3          10:05    Melbourne home  1       apple
4    visit 10:06    France                  apple

预期输出 #

source datetime location  type  start  company 
0          10:00    london    home  1       apple
1          10:03    unknown                 tesla
2    visit 10:04    France                  apple
3          10:05    Melbourne home  1       apple
4    visit 10:06    France                  apple

实现代码（Pandas） #

import pandas as pd

# 构造示例DataFrame
data = {
    'source': ['', '', '', '', 'visit'],
    'datetime': ['10:00', '10:03', '10:04', '10:05', '10:06'],
    'location': ['london', 'unknown', 'France', 'Melbourne', 'France'],
    'type': ['home', '', '', 'home', ''],
    'start': [1, '', '', 1, ''],
    'company': ['apple', 'tesla', 'apple', 'apple', 'apple']
}
df = pd.DataFrame(data)

# 1. 将datetime转为可计算的时间类型
df['datetime'] = pd.to_datetime(df['datetime'], format='%H:%M')

# 2. 按公司分组，提取每组内home行的时间并向前填充，得到每行最近的同公司home时间
df['last_home_time'] = df.groupby('company')['datetime'].where(df['type'] == 'home').ffill()

# 3. 计算当前行与最近home行的时间差（转成分钟）
df['time_diff_mins'] = (df['datetime'] - df['last_home_time']).dt.total_seconds() / 60

# 4. 更新source列：满足时间差≤5分钟且当前行不是home行时，标记为visit
df['source'] = df.apply(
    lambda row: 'visit' if (row['time_diff_mins'] <= 5 and row['type'] != 'home') else row['source'],
    axis=1
)

# 清理辅助列（可选）
df = df.drop(['last_home_time', 'time_diff_mins'], axis=1)

# 打印结果
print(df.to_string(index=False))

代码说明 #

1. 时间类型转换：将字符串格式的datetime转为Pandas时间类型，才能进行时间差计算。
2. 分组填充最近home时间：按公司分组后，用where筛选出home行的时间，再通过ffill()向前填充，确保每行都能拿到同公司最近的home记录时间。
3. 时间差计算：将时间差转为分钟，方便判断是否在5分钟范围内。
4. 更新source列：仅对非home行且满足时间差条件的记录，将source设为visit。

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