基于带阈值与操作符规则表的DataFrame过滤优化问询

需求说明 #

现有一个网格规则表，其行对应目标数据表的ColA区间，列对应ColB区间。需要给目标数据表新增一列：当ColA与ColB的数值组合落在规则表中值为1的区间时赋值1，否则赋值0。当前实现用了大量循环的复杂函数，寻求Pandas优化方案。

示例代码 #

import pandas as pd 
grids= pd.DataFrame(data = {"<0.1": [0,0,0,1], "<0.2": [0,0,1,1], ">=0.2": [0,0,1,0]}, index=["<0.1", "<0.2","<0.3", ">=0.3" ] ) 
df = pd.DataFrame({"ColA": [0.01, 0.15, 0.24, 0.5], "ColB": [0.15, 0.15, 0.24, 0.5]}) 

优化方案 #

核心思路是先给df的每个值匹配对应规则表的行/列标签，再通过索引直接从规则表中取值，全程用Pandas矢量化操作替代循环，效率和可读性都会提升。

步骤1：解析区间标签，匹配对应规则项

先把规则表的行/列字符串标签（比如<0.1、>=0.3）转换成可判断的逻辑，给df的ColA和ColB分别匹配对应的标签：

def match_interval(x, interval_labels):
    """根据给定的区间标签列表，返回x所属的区间标签"""
    for label in interval_labels:
        if '<' in label:
            threshold = float(label.split('<')[1])
            if x < threshold:
                return label
        elif '>=' in label:
            threshold = float(label.split('>=')[1])
            if x >= threshold:
                return label

# 给df添加匹配的规则标签列
df['ColA_rule'] = df['ColA'].apply(lambda x: match_interval(x, grids.index))
df['ColB_rule'] = df['ColB'].apply(lambda x: match_interval(x, grids.columns))

步骤2：从规则表中快速取值

把规则表转换成多层索引的Series，再通过df的标签列直接匹配取值，这一步完全避免循环：

# 将规则表转为MultiIndex格式的Series
grid_series = grids.stack()

# 用标签列作为索引匹配规则值，生成结果列
df['result'] = df.set_index(['ColA_rule', 'ColB_rule']).join(grid_series, how='left')[0]
df = df.reset_index(drop=True)

进阶优化：用pd.cut实现矢量化区间匹配

如果数据量很大，apply的效率还能再提升——用pd.cut直接做区间划分，完全矢量化：

# 解析ColA的区间阈值与标签
a_thresholds = []
a_labels = []
for label in grids.index:
    if '<' in label:
        val = float(label.split('<')[1])
        a_thresholds.append(val)
        a_labels.append(label)
    elif '>=' in label:
        val = float(label.split('>=')[1])
        a_thresholds.append(val)
        a_labels.append(label)
a_thresholds.sort()
a_bins = [-float('inf')] + a_thresholds

# 用pd.cut给ColA匹配规则标签
df['ColA_rule'] = pd.cut(df['ColA'], bins=a_bins, labels=a_labels, include_lowest=True)

# 同理处理ColB
b_thresholds = []
b_labels = []
for label in grids.columns:
    if '<' in label:
        val = float(label.split('<')[1])
        b_thresholds.append(val)
        b_labels.append(label)
    elif '>=' in label:
        val = float(label.split('>=')[1])
        b_thresholds.append(val)
        b_labels.append(label)
b_thresholds.sort()
b_bins = [-float('inf')] + b_thresholds

df['ColB_rule'] = pd.cut(df['ColB'], bins=b_bins, labels=b_labels, include_lowest=True)

# 再次用MultiIndex匹配取值
df['result'] = df.set_index(['ColA_rule', 'ColB_rule']).join(grid_series, how='left')[0]
df = df.reset_index(drop=True)

最终效果 #

运行后df会新增result列，值为1或0，完全符合需求，且代码简洁、效率远高于循环实现。

内容的提问来源于stack exchange，提问作者A.Papa