如何通过向量化加速Pandas DataFrame相关代码？求技术指导

问题描述 #

我是StackOverflow新手，目前使用Pandas导入CSV数据集进行处理，现有用于创建标签和数据偏差表征的函数，但代码加载耗时较长。由于对向量化技术不够熟悉，恳请告知如何通过向量化来加速这段代码。

原代码如下：

def create_labels(self):
    sensitive_label = {}
    for i in set(self.X_test[sensitive]):
        text = "Please Enter Label for Group" +" "+ str(i)+": "
        label = input(text)
        sensitive_label[i]=label
    return(sensitive_label)


def representation(self,sensitive, labels, predictions):
    full_table = self.X_test.copy()
    sens_df = {}
#Output is going to be a table
    for i in labels:
        full_table['p'] = predictions
        full_table['t'] = self.y_test
        sens_df[labels[i]] = full_table[full_table[sensitive]==i] #one table stored for female and one for male
    contigency_p = pd.crosstab(full_table[sensitive], full_table['t'])
    cp, pp, dofp, expectedp = chi2_contingency(contigency_p)
    contigency_pct_p = pd.crosstab(full_table[sensitive], full_table['t'], normalize='index')#p value of contigency table
    sens_rep = {}
    for i in labels:
        sens_rep[labels[i]] = (self.X_test[sensitive].value_counts()/self.X_test[sensitive].value_counts().sum())[i]
    labl_rep = {}
    for i in labels:
        labl_rep[str(i)] = (self.y_test.value_counts()/self.y_test.value_counts().sum())[i]
    fig = make_subplots(rows=1, cols=2)
    for i in labels:
        fig.add_trace(go.Bar(
        showlegend=False,
        x = [labels[i]],
        y= [sens_rep[labels[i]]]),row=1,col=1)
        fig.add_trace(go.Bar(
        showlegend=False,
        x = [str(i)],
        y= [labl_rep[str(i)]],
        marker_color=['orange','blue'][i]),row=1,col=2)
    c, p, dof, expected = chi2_contingency(contigency_p)
    cont_table = (tabulate(contigency_pct_p.T, headers=labels.values(), tablefmt='fancy_grid'))
    return cont_table, sens_df, fig, p
        #sens_df dataset based on the senstive labels

优化方案：用向量化操作替代循环，减少冗余计算 #

你的代码慢的核心原因是多次循环重复计算Pandas Series，以及不必要的重复赋值/统计。以下是针对性优化：

1. 优化create_labels函数

用Pandas原生的unique()替代set()，避免数据类型转换开销，尤其在大数据集下更高效：

def create_labels(self):
    sensitive_label = {}
    # 用unique()替代set()，保留原数据顺序且更高效
    for i in self.X_test[sensitive].unique():
        text = f"Please Enter Label for Group {str(i)}: "
        label = input(text)
        sensitive_label[i] = label
    return sensitive_label

2. 大幅优化representation函数

重点解决循环重复计算、冗余操作的问题：

def representation(self, sensitive, labels, predictions):
    full_table = self.X_test.copy()
    # 把p和t的赋值提到循环外，避免重复执行
    full_table['p'] = predictions
    full_table['t'] = self.y_test
    
    # 用groupby替代循环筛选，一次分组生成所有子表
    sens_df = {labels[group_key]: group_df for group_key, group_df in full_table.groupby(sensitive)}
    
    # 仅计算一次卡方检验，避免重复计算
    contigency_p = pd.crosstab(full_table[sensitive], full_table['t'])
    c, p, dof, expected = chi2_contingency(contigency_p)
    contigency_pct_p = pd.crosstab(full_table[sensitive], full_table['t'], normalize='index')
    
    # 向量化计算敏感组占比，仅计算一次value_counts
    sens_counts = self.X_test[sensitive].value_counts(normalize=True)
    sens_rep = {labels[k]: sens_counts[k] for k in labels}
    
    # 向量化计算标签占比，仅计算一次value_counts
    labl_counts = self.y_test.value_counts(normalize=True)
    labl_rep = {str(k): labl_counts[k] for k in labels}
    
    # 图表绘制：循环逻辑不变，但数据已提前用向量化准备好
    fig = make_subplots(rows=1, cols=2)
    for group_name, ratio in sens_rep.items():
        fig.add_trace(go.Bar(x=[group_name], y=[ratio], showlegend=False), row=1, col=1)
    for k in labels:
        fig.add_trace(go.Bar(
            x=[str(k)], 
            y=[labl_rep[str(k)]], 
            marker_color=['orange','blue'][k], 
            showlegend=False
        ), row=1, col=2)
    
    cont_table = tabulate(contigency_pct_p.T, headers=labels.values(), tablefmt='fancy_grid')
    return cont_table, sens_df, fig, p

优化点说明 #

• 去除冗余计算：原代码两次调用chi2_contingency、循环内重复赋值full_table['p']和full_table['t']，优化后仅执行一次。
• 向量化统计：把循环内多次计算的value_counts()提到循环外，一次计算完成所有比例映射，避免重复遍历数据集。
• 高效分组：用groupby替代循环筛选子表，Pandas的groupby内部是向量化实现，比手动循环切片快得多。

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