基于分组与概率区间匹配规则合并两个DataFrame的技术问询

Let's break this down with concrete examples and actionable solutions. First, let's set up sample DataFrames to mirror your scenario clearly:

Sample Data #

import pandas as pd

# df1: Contains groups and their corresponding probability interval thresholds
df1 = pd.DataFrame({
    'group': ['A', 'A', 'B', 'B'],
    'prob_interval': [0.3, 0.5, 0.2, 0.7]
})

# df2: Contains groups and random numbers to map
df2 = pd.DataFrame({
    'group': ['A', 'A', 'B', 'B'],
    'random_num': [0.2, 0.4, 0.1, 0.5]
})

In this example:

• For group A, 0.2 maps to 0.3 (smallest interval greater than 0.2), and 0.4 maps to 0.5
• For group B, 0.1 maps to 0.2, and 0.5 maps to 0.7

Method 1: Use merge_asof (Efficient for Large Datasets) #

pandas.merge_asof is perfect for this scenario because it can match values to the nearest key in a sorted dataset, with control over matching direction. Here's how to use it:

Step 1: Sort Both DataFrames #

merge_asof requires the matching columns to be sorted, so we'll sort by group and the relevant numeric column:

# Sort df1 by group and prob_interval (ascending)
df1_sorted = df1.sort_values(by=['group', 'prob_interval']).reset_index(drop=True)

# Sort df2 by group and random_num (ascending)
df2_sorted = df2.sort_values(by=['group', 'random_num']).reset_index(drop=True)

Step 2: Perform Forward Match #

Use direction='forward' to find the smallest interval greater than or equal to the random number, grouped by the group column:

result = pd.merge_asof(
    df2_sorted,
    df1_sorted,
    left_on='random_num',
    right_on='prob_interval',
    by='group',
    direction='forward'
)

# Clean up and reorder columns if needed
result = result[['group', 'random_num', 'prob_interval']].rename(columns={'prob_interval': 'mapped_interval'})

Output: #

Method 2: Groupby + searchsorted (Flexible for Custom Logic) #

If you need more control (e.g., handling random numbers larger than all intervals), use groupby combined with numpy.searchsorted:

Step 1: Prepare Interval Lookup #

First, create a dictionary mapping each group to its sorted list of intervals:

import numpy as np

# Get sorted intervals per group
interval_lookup = df1.groupby('group')['prob_interval'].apply(lambda x: np.sort(x)).to_dict()

Step 2: Define Mapping Function #

Create a function to map random numbers for each group, with custom handling for out-of-range values:

def map_random_values(group_df):
    group_name = group_df.name
    intervals = interval_lookup[group_name]
    
    # Find indices of the smallest interval greater than each random number
    indices = intervals.searchsorted(group_df['random_num'], side='right')
    
    # Map indices to interval values; handle numbers larger than all intervals
    group_df['mapped_interval'] = [
        intervals[i] if i < len(intervals) else intervals[-1]
        for i in indices
    ]
    return group_df

# Apply the function to each group in df2
result = df2.groupby('group', group_keys=False).apply(map_random_values)

This method lets you easily adjust behavior (e.g., replace out-of-range values with NaN instead of the maximum interval).

Key Notes #

• Ensure your probability intervals are unique per group (if duplicates exist, merge_asof will pick the first matching value).
• For very large datasets, merge_asof is significantly faster than groupby+apply due to vectorized operations.

内容的提问来源于stack exchange，提问作者G. Deg