如何基于指定中位数与尾点数从DataFrame生成数据区间？

我有一个包含受试者信息的大型DataFrame，以下是简化后的示例：

subject   age   sex   
A         5.35  Female
B         5.70  Male
C         6.00  Female
D         6.07  Male
E         6.25  Male
F         6.88  Male
G         7.00  Female
H         7.02  Male
I         7.11  Female
J         8.00  Male
K         8.50  Female

需要编写一个函数，输入参数为age（作为区间中位数）、tail（中位数两侧需要选取的数据点数量），后续还要扩展支持sex参数。期望输出从原DataFrame中筛选出的区间DataFrame，且df.loc无法直接满足需求。示例功能如下：

def interval_set(age = 7, tail = 3):
   # 处理上述DataFrame后输出如下结果

Out[1]:

subject   age   sex   
D         6.07  Male
E         6.25  Male
F         6.88  Male
G         7.00  Female
H         7.02  Male
I         7.11  Female
J         8.00  Male

核心思路是先定位到与目标中位数最接近的行，再向两侧各取tail个数据点，同时处理边界避免索引越界。

基础实现代码 #

import pandas as pd

# 构造示例DataFrame
df = pd.DataFrame({
    'subject': ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K'],
    'age': [5.35, 5.70, 6.00, 6.07, 6.25, 6.88, 7.00, 7.02, 7.11, 8.00, 8.50],
    'sex': ['Female', 'Male', 'Female', 'Male', 'Male', 'Male', 'Female', 'Male', 'Female', 'Male', 'Female']
})

def interval_set(age=7, tail=3):
    # 1. 找到与目标age最接近的行索引（存在多个相同值时取第一个）
    target_idx = (df['age'] - age).abs().idxmin()
    # 2. 计算安全的起始/结束索引，避免越界
    start_idx = max(0, target_idx - tail)
    end_idx = min(len(df)-1, target_idx + tail)
    # 3. 截取区间并返回副本（防止修改原DataFrame）
    return df.loc[start_idx:end_idx].copy()

# 测试示例
print(interval_set(7, 3))

代码说明 #

• 用绝对差最小值定位目标中位数对应的行，确保精准找到区间中心；
• 通过max(0, ...)和min(len(df)-1, ...)处理边界情况，比如当目标行靠近DataFrame开头/结尾时，不会出现无效索引；
• 返回副本是为了避免后续操作影响原数据。

扩展支持sex参数的版本 #

如果需要按性别筛选后再取区间，修改函数如下：

def interval_set(age=7, tail=3, sex=None):
    # 先按性别过滤（未指定则用原数据）
    filtered_df = df.copy()
    if sex is not None:
        filtered_df = filtered_df[filtered_df['sex'] == sex]
    
    # 后续逻辑同基础版本
    target_idx = (filtered_df['age'] - age).abs().idxmin()
    start_idx = max(0, target_idx - tail)
    end_idx = min(len(filtered_df)-1, target_idx + tail)
    return filtered_df.loc[start_idx:end_idx].copy()

# 示例：筛选女性中以7为中位数、两侧各2个数据点的区间
print(interval_set(7, 2, sex='Female'))

内容的提问来源于stack exchange，提问作者florence-y