收容所每日在住动物数计算代码的性能优化求助

问题背景 #

我使用当地动物收容所的动物入出记录数据集，核心字段为intake_date（入所日期）和outcome_date（出所日期），需要计算每日在住动物数（当晚留宿的动物）并绘制变化曲线，计算规则如下：

• 若outcome_date为空，该动物从intake_date到当前日期的每一天都计入统计；
• 若outcome_date不为空，该动物从intake_date到outcome_date前一天的每一天都计入统计。

原实现代码采用嵌套循环遍历日期和每条记录，面对25000+条记录、超10年的时间跨度，循环执行次数超过9000万次，耗时数小时，急需性能优化。

原代码：

import pandas as pd
import time
import datetime


def fetch_data():
    url = "https://raw.githubusercontent.com/Etharialle/shelter-streamlit/main/datasets/asio.csv"
    df = pd.read_csv(url, on_bad_lines='skip').replace("'", "", regex=True)
    return df


df = fetch_data()
# Ensure 'Intake Date' is a datetime type
df_open_cases["Intake Date"] = pd.to_datetime(
    df_open_cases["Intake Date"])
df_open_cases["Outcome Date"] = pd.to_datetime(
    df_open_cases["Outcome Date"])

# Animals in Residence by Date
start_date = df_open_cases["Intake Date"].min()
end_date = datetime.datetime.now()
date_list = {}

while start_date <= end_date:
    loop_date = start_date.strftime('%Y-%m-%d')
    date_list[loop_date] = 0
    start_date += datetime.timedelta(days=1)

for k, v, in date_list.items():
    for index, rows in df_open_cases.iterrows():
        if pd.isnull(rows['Outcome Date']) and rows["Intake Date"].strftime('%Y-%m-%d') <= k:
            date_list[k] += 1
        elif rows["Intake Date"].strftime('%Y-%m-%d') <= k and rows['Outcome Date'].strftime('%Y-%m-%d') > k:
            date_list[k] += 1
        else:
            pass
df_dates = pd.DataFrame.from_dict(
    date_list, orient='index', columns=['Animals in Residence'])

优化方案：事件计数法（向量化处理） #

放弃嵌套循环，改用事件累加的思路，将每个动物的入出转化为日期节点的计数变化，再通过累加得到每日在住数，全程用Pandas向量化操作，效率提升几个数量级。

步骤说明 #

1. 数据预处理：
   • 将日期字段转为datetime类型
   • 把空的outcome_date填充为当前日期（因为无出所记录的动物需统计到当前）
2. 生成事件数据：
   • 每条记录对应两个事件：入所日期+1，出所日期-1（出所当天不计入统计，所以出所日期就是停止计数的节点，当天减1）
3. 聚合与累加：
   • 按日期对事件求和，得到每日的计数变化量
   • 对变化量做累加，得到每日的在住动物总数
4. 补全连续日期：
   • 生成从最早入所日期到当前日期的完整日期序列，确保没有缺失日期，缺失日期的变化量为0

优化后代码 #

import pandas as pd
import datetime

def fetch_data():
    url = "https://raw.githubusercontent.com/Etharialle/shelter-streamlit/main/datasets/asio.csv"
    df = pd.read_csv(url, on_bad_lines='skip').replace("'", "", regex=True)
    return df

# 加载并预处理数据
df = fetch_data()
df['Intake Date'] = pd.to_datetime(df['Intake Date'])
df['Outcome Date'] = pd.to_datetime(df['Outcome Date'])

# 填充空的出所日期为当前日期
current_date = datetime.datetime.now().date()
df['Outcome Date'] = df['Outcome Date'].fillna(pd.Timestamp(current_date))

# 生成入所事件（+1）和出所事件（-1）
intake_events = df[['Intake Date']].assign(count=1).rename(columns={'Intake Date': 'date'})
outcome_events = df[['Outcome Date']].assign(count=-1).rename(columns={'Outcome Date': 'date'})

# 合并事件并按日期求和
all_events = pd.concat([intake_events, outcome_events])
daily_changes = all_events.groupby('date')['count'].sum().reset_index()

# 生成完整的日期序列
start_date = df['Intake Date'].min().date()
end_date = current_date
full_dates = pd.date_range(start=start_date, end=end_date, freq='D').to_frame(name='date')

# 合并完整日期与每日变化，填充缺失值为0
daily_changes = pd.merge(full_dates, daily_changes, on='date', how='left').fillna(0)

# 计算累加和得到每日在住动物数
daily_changes['Animals in Residence'] = daily_changes['count'].cumsum()

# 整理结果格式
result_df = daily_changes[['date', 'Animals in Residence']].set_index('date')
print(result_df.head())

效果说明 #

此方案将时间复杂度从O(N*D)（N为记录数，D为天数）降至O(N+D)，对于25000条记录+3650天的场景，执行时间可从数小时缩短至几秒内，完全满足性能需求。

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