Python递归实现代码优化：以数据分类存储为例

在处理数据时，我们经常需要根据不同的条件对数据进行分类存储。如果分类层级较多，使用嵌套循环会导致代码冗长、难以维护。本文将介绍如何使用递归函数优化这类代码，并以一个实际案例进行说明。

问题背景

假设我们需要从一个数据源中提取数据，并根据数据的类型和类别进行分类存储。数据源、类型和类别的配置信息存储在一个名为 config 的对象中。

以下是用嵌套循环实现的代码：

def process_data(config, save_filePath, now_day, save_dict):
    for tj_config in config.tj_config.keys():
        tj_config_data = self.pickUpType(self.getData(),tj_config)   #按type提取数据
        save_typefilename = tj_config +'_'+ now_day + '.csv'
        save_typefile = os.path.join(save_filePath,save_typefilename)    #该类的保存文件名
        TianJiIocs.save_Csv(tj_config_data, save_typefile)  # 生成Type文件
        #保存数据到字典
        save_dict[tj_config] = {}  # json
        save_dict[tj_config]['dataFile'] = save_typefile
        save_dict[tj_config]['count'] = f'{tj_config}类-数据总量:{tj_config_data.shape[0]}'

        #威胁类
        category_config = config.tj_config[tj_config]['category']

        for category in category_config.keys():
            category_data = self.pickUpCategory(tj_config_data,category)
            save_categoryFilename = tj_config + '_' + category + '_' + now_day + '.csv'
            save_categoryFile = os.path.join(save_filePath, save_categoryFilename)  # 该威胁类型的保存文件名
            TianJiIocs.save_Csv(category_data, save_categoryFile)  # 生成威胁类文件

递归优化

我们可以使用递归函数将嵌套循环逻辑简化，提高代码的可读性和可维护性。

def recursive_function(config, save_filePath, now_day, save_dict, data=None, tj_config=None):
    if data is None:
        # 初始调用，处理所有类型
        for tj_config in config.tj_config.keys():
            recursive_function(config, save_filePath, now_day, save_dict, self.getData(), tj_config)
        return

    tj_config_data = self.pickUpType(data, tj_config)  # 按type提取数据
    save_typefilename = tj_config + '_' + now_day + '.csv'
    save_typefile = os.path.join(save_filePath, save_typefilename)  # 该类的保存文件名
    TianJiIocs.save_Csv(tj_config_data, save_typefile)  # 生成Type文件
    # 保存数据到字典
    save_dict[tj_config] = {}
    save_dict[tj_config]['dataFile'] = save_typefile
    save_dict[tj_config]['count'] = f'{tj_config}类-数据总量:{tj_config_data.shape[0]}'

    # 威胁类
    category_config = config.tj_config[tj_config]['category']
    for category in category_config.keys():
        # 递归调用，处理每个类别
        recursive_function(config, save_filePath, now_day, save_dict, tj_config_data, category)

代码解释

recursive_function 函数接受 config, save_filePath, now_day, save_dict, data, tj_config 参数。其中 data 和 tj_config 用于递归调用时传递数据和当前处理的类型。
当 data 为 None 时，表示是初始调用，遍历所有类型并递归调用自身处理每个类型的数据。
提取当前类型的数据，保存到文件，并更新 save_dict。
遍历当前类型的类别，递归调用自身处理每个类别的数据。

总结

使用递归函数优化嵌套循环逻辑可以使代码更简洁、易读、易维护。在处理多层级数据分类存储时，递归是一种有效的解决方案。