Pandas使用正则匹配可变数量列 按公式批量计算生成新列

示例源数据 #

需求规则 #

使用正则表达式regex = '^p\d'匹配符合命名规则的列，自动生成n1、n2...格式的新列，计算逻辑：

• 逐行计算所有匹配到的p列的平方和的平方根作为公共分母
• 每个n列值 = 对应序号p列的值 / 该行公共分母
• 自动适配不同CSV文件：p列位置不固定、p列总数不固定时，都能自动匹配所有目标列，无匹配p列时自动停止生成，无需手动指定列参数

初始测试代码如下：

import pandas as pd
import numpy as np

data_1 = {'p1': [1, 2, 3],
          'p2': [4, 5, 6],
          'p3': [3, 70 ,6],
          'p4': [40, 8, 500],
         }

df1 = pd.DataFrame(data = data_1)
df1

实现代码 #

直接用pandas内置的正则筛选方法获取目标列，无需硬编码列名或列数：

# 按指定正则匹配所有p开头的目标列，自动适配列位置和列总数
p_cols = df1.filter(regex=r'^p\d', axis=1).columns

# 逐行计算所有p列平方和的平方根，作为归一化分母
row_l2_norm = np.sqrt((df1[p_cols] ** 2).sum(axis=1))

# 遍历匹配到的p列，按顺序生成对应n列，遍历完所有p列自动停止
for col_idx, p_col in enumerate(p_cols, start=1):
    df1[f'n{col_idx}'] = df1[p_col] / row_l2_norm

运行结果 #

处理后DataFrame输出如下：

逻辑说明：filter方法会按正则自动扫描所有列名，不管p列在表中什么位置、总共有多少个p列都能准确命中；遍历过程随匹配到的p列数量自动终止，不需要提前判断列数，适配任意结构的同类CSV文件。

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