求助：将pandas IntervalIndex作为MultiIndex一部分的匹配实现问题

Hey there, let's work through this IntervalIndex + MultiIndex matching problem together. I've dealt with similar headaches before, so I know exactly where the sticking points are. First, let's recap your setup to make sure we're on the same page, then walk through two reliable solutions.

Step 1: Reproduce Your Data #

First, let's recreate your df1 and define a sample df2 (since you didn't share its full data, I'll use a realistic example that aligns with your description):

import pandas as pd

# Build df1 exactly as you described
data1 = {
    'month': [1,1,1,1,1,1,1,1,2,2,2,2],
    'group': ['A','A','A','B','B','B','C','C','A','A','B','B'],
    'time': [10,20,25,5,7,17,30,35,5,10,17,19],
    'distance': [100,120,110,90,99,60,200,170,40,27,33,65]
}
df1 = pd.DataFrame(data1)
df1.set_index(['month', 'group', 'time'], inplace=True)

# Sample df2 with month, group, start, end columns
data2 = {
    'month': [1,1,2,2],
    'group': ['A','B','A','B'],
    'start': [15,5,0,15],
    'end': [25,15,10,20]
}
df2 = pd.DataFrame(data2)

Step 2: Solution 1: Grouped Interval Matching (Efficient for Large Data) #

This approach groups data by month and group first, then matches time values to intervals within each group. It's great if you're working with large datasets since it avoids a full cross-merge.

# Reset df1's index to work with columns instead of MultiIndex
df1_reset = df1.reset_index()

def match_group_intervals(group_df, interval_df):
    # Grab the current group's month and group label
    current_month = group_df['month'].iloc[0]
    current_group = group_df['group'].iloc[0]
    
    # Filter df2 to only intervals for this (month, group) pair
    group_intervals = interval_df.xs((current_month, current_group), level=['month', 'group'], drop_level=False)
    if group_intervals.empty:
        return pd.DataFrame()  # Return empty if no intervals match this group
    
    # Create an IntervalIndex for the current group's intervals
    interval_idx = pd.IntervalIndex.from_arrays(group_intervals['start'], group_intervals['end'], closed='left')
    
    # Match each time value to the intervals that contain it
    group_df['matched_interval'] = group_df['time'].apply(
        lambda t: interval_idx[interval_idx.contains(t)].tolist() if any(interval_idx.contains(t)) else pd.NA
    )
    
    # Expand rows where a time matches multiple intervals, drop unmatched rows
    group_df = group_df.explode('matched_interval').dropna(subset=['matched_interval'])
    
    # Merge back in start/end values for clarity
    group_df = group_df.merge(
        group_intervals.reset_index()[['start', 'end']],
        left_on=group_df['matched_interval'].apply(lambda iv: (iv.left, iv.right)),
        right_on=group_intervals.reset_index()[['start', 'end']].apply(tuple, axis=1),
        how='left'
    )
    
    # Reconstruct the MultiIndex with interval info
    group_df.set_index(['month', 'group', 'time', 'matched_interval'], inplace=True)
    group_df.index.rename(['month', 'group', 'time', 'time_interval'], inplace=True)
    return group_df

# Apply the function to each (month, group) group and concatenate results
final_result = pd.concat([
    match_group_intervals(g, df2.set_index(['month', 'group'])) 
    for _, g in df1_reset.groupby(['month', 'group'])
])

print(final_result)

Step 3: Solution 2: Cross-Merge + Filter (Simple for Smaller Data) #

If your dataset isn't massive, this more straightforward method works well. We cross-merge on month and group, then filter rows where time falls within the interval.

# Reset indexes for both DataFrames
df1_reset = df1.reset_index()
df2_reset = df2.reset_index()

# Merge on shared keys (month, group) first
cross_merged = df1_reset.merge(df2_reset, on=['month', 'group'], how='inner')

# Filter rows where time is within [start, end) (adjust closed logic if needed)
matched_rows = cross_merged[cross_merged.apply(lambda row: row['start'] <= row['time'] < row['end'], axis=1)]

# Create IntervalIndex and rebuild the MultiIndex
matched_rows['time_interval'] = pd.IntervalIndex.from_arrays(matched_rows['start'], matched_rows['end'], closed='left')
matched_rows.set_index(['month', 'group', 'time', 'time_interval'], inplace=True)

print(matched_rows)

Key Notes to Avoid Pitfalls #

• Interval Closure: The closed='left' parameter means intervals are [start, end). Adjust this to 'right', 'both', or 'neither' based on your exact matching needs.
• Multiple Matches: If a single time value falls into multiple intervals, the explode method (in Solution 1) will create separate rows for each match.
• Unmatched Rows: Both solutions drop rows where no interval matches. If you want to retain these, remove the dropna call (Solution 1) or switch to a left merge and handle NaN values (Solution 2).

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