基于日期匹配为Pandas DataFrame添加列并处理日期重叠

Got it, let's tackle this problem step by step. The goal is to add a plan column to df1 based on date range matches from df2, with NaN assigned to any dates that overlap multiple ranges. Here's a robust, easy-to-follow solution:

完整函数实现（基础版，适合中小数据集） #

This version uses pandas' cross merge and grouping logic, which is intuitive and easy to debug:

import pandas as pd

def add_plan_to_df1(df1, df2):
    # Create copies to avoid modifying original DataFrames
    df1_copy = df1.copy()
    df2_copy = df2.copy()
    
    # Convert all date columns to datetime type (critical for accurate range checks)
    df1_copy['Date'] = pd.to_datetime(df1_copy['Date'])
    df2_copy['From'] = pd.to_datetime(df2_copy['From'])
    df2_copy['to'] = pd.to_datetime(df2_copy['to'])
    
    # Filter out invalid ranges where From > to (prevents errors)
    df2_copy = df2_copy[df2_copy['From'] <= df2_copy['to']].reset_index(drop=True)
    
    # Cross merge to generate all possible date-range combinations, then filter valid matches
    merged = pd.merge(df1_copy, df2_copy, how='cross')
    valid_matches = merged[(merged['Date'] >= merged['From']) & (merged['Date'] <= merged['to'])]
    
    # Count how many plans each date matches (to detect overlaps)
    match_counts = valid_matches.groupby('Date')['plan'].count().reset_index(name='overlap_count')
    
    # Merge valid plans back to df1, keeping only one plan per date (if no overlap)
    df1_with_plans = pd.merge(
        df1_copy,
        valid_matches[['Date', 'plan']].drop_duplicates('Date'),
        on='Date',
        how='left'
    )
    
    # Assign NaN to dates with overlapping ranges, or no matching ranges
    df1_with_plans = pd.merge(df1_with_plans, match_counts, on='Date', how='left')
    df1_with_plans['plan'] = df1_with_plans.apply(
        lambda row: pd.NA if (pd.isna(row['overlap_count']) or row['overlap_count'] > 1) else row['plan'],
        axis=1
    )
    
    # Clean up and return the original column order
    return df1_with_plans[['Date', 't_factor', 'plan']]

优化版（适合大数据量） #

If you're working with millions of rows, cross merge can be memory-heavy. This version uses numpy broadcasting for faster range checks:

import pandas as pd
import numpy as np

def add_plan_to_df1(df1, df2):
    df1_copy = df1.copy()
    df2_copy = df2.copy()
    
    # Convert dates and clean invalid ranges
    df1_copy['Date'] = pd.to_datetime(df1_copy['Date'])
    df2_copy['From'] = pd.to_datetime(df2_copy['From'])
    df2_copy['to'] = pd.to_datetime(df2_copy['to'])
    df2_copy = df2_copy[df2_copy['From'] <= df2_copy['to']].reset_index(drop=True)
    
    # Convert dates to numpy arrays for broadcasting
    dates = df1_copy['Date'].values[:, np.newaxis]
    from_dates = df2_copy['From'].values[np.newaxis, :]
    to_dates = df2_copy['to'].values[np.newaxis, :]
    
    # Create a mask where each cell indicates if a date falls in a range
    match_mask = (dates >= from_dates) & (dates <= to_dates)
    
    # Calculate overlap counts and get first matching plan
    overlap_counts = match_mask.sum(axis=1)
    first_match_idx = match_mask.argmax(axis=1)
    first_match_idx[overlap_counts == 0] = -1  # Mark no matches
    
    # Assign plan values, with NaN for overlaps or no matches
    df1_copy['plan'] = np.where(
        (overlap_counts == 1),
        df2_copy['plan'].values[first_match_idx],
        pd.NA
    )
    
    return df1_copy[['Date', 't_factor', 'plan']]

测试用例验证 #

Let's test with sample data to confirm the logic works:

# Sample input DataFrames
df1 = pd.DataFrame({
    'Date': ['2024-01-01', '2024-01-05', '2024-01-10', '2024-01-15'],
    't_factor': [1.2, 3.4, 5.6, 7.8]
})

df2 = pd.DataFrame({
    'From': ['2024-01-01', '2024-01-08'],
    'to': ['2024-01-07', '2024-01-12'],
    'plan': ['A', 'B'],
    'score': [90, 85]
})

# Run the function
result = add_plan_to_df1(df1, df2)
print(result)

Expected Output: #

Date  t_factor  plan
0 2024-01-01       1.2     A
1 2024-01-05       3.4     A
2 2024-01-10       5.6  <NA>
3 2024-01-15       7.8  <NA>

关键细节说明 #

• Date Type Conversion: Always convert date columns to datetime to avoid string comparison errors.
• Invalid Range Handling: We filter out rows in df2 where From > to to prevent nonsensical range checks.
• Overlap Detection: Any date matching more than one range gets assigned NaN, as required.
• No Side Effects: The function works on copies of the input DataFrames, so your original data remains untouched.

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