基于日期计算Pandas DataFrame各ID平均销售额的优化方案

Great question! Using np.select works for small ranges, but it's definitely not scalable when you have 12+ months. Here are two much cleaner, efficient approaches that will work no matter how many months you add:

Approach 1: Vectorized Column-Wise Calculation (Fast & Direct) #

This method avoids reshaping your data and uses vectorized operations to compute the average in one go—perfect for large datasets where speed matters.

import pandas as pd
import numpy as np

# Your original data setup
a = [1, 2, 3, 4, 5]
b = ['2019-08-01', '2019-09-01', '2019-10-23', '2019-11-12', '2019-11-30']
c = [12, 0, 0, 0, 0]
d = [0, 23, 0, 0, 0]
e = [12, 24, 35, 0, 0]
f = [0, 0, 44, 56, 82]
g = [21, 22, 17, 75, 63]
df = pd.DataFrame({'ID': a, 'Date': b, 'Unit_sold_8': c, 'Unit_sold_9': d, 'Unit_sold_10': e, 'Unit_sold_11': f, 'Unit_sold_12': g})
df['Date'] = pd.to_datetime(df['Date'])

# 1. Extract the start month from each ID's activation date
df['start_month'] = df['Date'].dt.month

# 2. Isolate sales columns and get their corresponding month numbers
sales_cols = [col for col in df.columns if col.startswith('Unit_sold_')]
month_numbers = np.array([int(col.split('_')[-1]) for col in sales_cols])

# 3. Create a mask to filter valid sales months (>= start month)
start_month_matrix = df['start_month'].values[:, np.newaxis]
valid_sales_mask = month_numbers >= start_month_matrix

# 4. Calculate total valid sales and count of valid months
total_sales = (df[sales_cols].values * valid_sales_mask).sum(axis=1)
valid_month_count = valid_sales_mask.sum(axis=1)

# 5. Compute average sales
df['average_sales'] = total_sales / valid_month_count

# Clean up if needed
df = df.drop('start_month', axis=1)

Approach 2: Reshape to Long Format (Readable & Flexible) #

This method reshapes your data into a "long" format, making it easier to filter and aggregate. It’s ideal if you might need to compute other metrics (like sum or median) later without rewriting core logic.

# Use the same original df as above

# 1. Extract start month
df['start_month'] = df['Date'].dt.month

# 2. Reshape wide data to long format
melted_df = pd.melt(
    df,
    id_vars=['ID', 'Date', 'start_month'],
    var_name='sales_month_col',
    value_name='unit_sold'
)

# 3. Get numeric month from the sales column name
melted_df['sales_month'] = melted_df['sales_month_col'].str.extract(r'(\d+)').astype(int)

# 4. Filter rows where sales month is on or after activation month
valid_sales = melted_df[melted_df['sales_month'] >= melted_df['start_month']]

# 5. Calculate average sales per ID
average_sales = valid_sales.groupby('ID')['unit_sold'].mean().reset_index(name='average_sales')

# 6. Merge results back to original DataFrame
df = df.merge(average_sales, on='ID').drop('start_month', axis=1)

Key Advantages Over np.select #

• Scalability: Both methods work for any number of months (12, 24, etc.) without adding more code.
• Maintainability: No need to write dozens of condition-value pairs for each month.
• Efficiency: Vectorized operations (Approach 1) and pandas' optimized groupby (Approach 2) are far faster than row-wise operations or long np.select chains.

Content of the question来源于stack exchange，提问作者Long_NgV