Conditional grouped cumulative sum实现患者就诊指标高效计算问询

你原来的逐行循环方案时间复杂度为O(n²)（每行都要遍历同组所有历史记录），百万行规模下耗时极长是必然的。改用向量化运算/数据库窗口函数可以把时间复杂度降到O(n log n)，普通配置下百万行数据跑完只需数分钟。
以下是不同场景下的实现方案：

1. 数据库SQL实现（适合数据存储在关系型数据库的场景） #

直接用窗口函数实现，无需额外拉取数据到本地，假设表名为patient_visits：

SELECT 
    id,
    visit_date,
    -- 总既往就诊次数：统计同id下早于当前就诊日期的记录总数
    COUNT(*) OVER (PARTITION BY id ORDER BY visit_date ROWS BETWEEN UNBOUNDED PRECEDING AND 1 PRECEDING) AS prior_visits,
    -- 近1年既往就诊次数：统计同id下在当前就诊日期前1年内的记录数
    COUNT(CASE WHEN visit_date >= DATE_SUB(visit_date, INTERVAL 1 YEAR) THEN 1 END) OVER (PARTITION BY id ORDER BY UNIX_TIMESTAMP(visit_date) RANGE BETWEEN 31536000 PRECEDING AND 1 PRECEDING) AS prior_one_year_visits
FROM patient_visits
ORDER BY id, visit_date;

注：不同数据库语法略有差异，比如PostgreSQL可直接用RANGE BETWEEN INTERVAL '1 year' PRECEDING AND INTERVAL '0 days' PRECEDING简化日期范围判断。

2. Python Pandas实现（适合单机可容纳的百万级数据集） #

用分组窗口替代循环，性能提升百倍以上：

import pandas as pd

# 先将就诊日期转为日期格式
df['visit_date'] = pd.to_datetime(df['visit_date'])
# 按患者id、就诊日期升序排序
df = df.sort_values(['id', 'visit_date']).reset_index(drop=True)

# 计算总既往就诊次数
df['prior_visits'] = df.groupby('id').cumcount()

# 计算1年内既往就诊次数：滚动窗口设为365天，左闭右开排除当前行日期
def count_recent_one_year(x):
    return x.rolling('365D', on='visit_date', closed='left').count()['id']
df['prior_one_year_visits'] = df.groupby('id', group_keys=False).apply(count_recent_one_year)

如果要进一步优化性能，可改用numpy向量化条件判断替代rolling方法。

3. PySpark实现（适合超大规模数据集，超出单机内存的场景） #

逻辑和SQL窗口函数一致，支持分布式并行计算：

from pyspark.sql import Window
from pyspark.sql.functions import col, count, expr

# 总就诊次数窗口规则
w_total = Window.partitionBy('id').orderBy('visit_date').rowsBetween(Window.unboundedPreceding, -1)
# 近1年就诊次数窗口规则，按时间戳换算1年时长
w_one_year = Window.partitionBy('id').orderBy(expr("unix_timestamp(visit_date)")).rangeBetween(-365*24*3600, -1)

df = df.withColumn('prior_visits', count('*').over(w_total)) \
       .withColumn('prior_one_year_visits', count('*').over(w_one_year))

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