Python: Pandas - 嵌套循环需要很长时间才能完成。如何加快速度？

Posted 2023-03-12

【中文标题】Python: Pandas - 嵌套循环需要很长时间才能完成。如何加快速度？

【英文标题】：Python: Pandas - Nested Loop takes very long to complete. How to speed up?

【发布时间】：2021-08-15 17:20:56

【问题描述】：

我正在自学 Python 和 Pandas 以支持我的日常工作。经过大量试验和错误，我构建了以下功能。该函数将 (i) 引用为“数据集”的数据框、(ii) 国家名称列表和 (iii) 唯一法人实体 ID 列表作为参数。 （该功能有效。）

数据集是一个包含 300,000 多行和大约 30 列的大型数据框——它是总账的转储。关键列是“LE_ID”和“COUNTRY”，它们分别包含 (i) 相关法人实体的唯一 ID 和 (ii) 该法人实体所在国家/地区的名称。并非所有行都是唯一的，大约有 5000 个 LE_ID 填充了 300,000 多行。

我想将此数据集“拆分”为 XLS 文件，这些文件按国家、大纲、每个选项卡、每个 LE_ID 的分类帐详细信息。下面的函数实现了这一点。但是速度太快了——在我最近的笔记本电脑上完成需要 40 分钟。

功能：

def SplitDatasetByCountry(dataset, countries, ids):
        
    for country in countries:

        ## output
        output_folder = Path('/Users/XXXXXX/Desktop/TOOL/Reports')
        output_filename = 'Data__for_' + str(country) + '_.xlsx'
        output = output_folder / output_filename
    
        ## writer
        writer = pd.ExcelWriter(output)
        workbook = writer.book

        ## country logic
        x = dataset.loc[dataset['COUNTRY'] == country]
        
        ids_for_entities_in_country = x['LE_ID'].to_list()
        unique_ids = list(set(ids_for_entities_in_country))

        for id in ids:

            if id in unique_ids:
                y = x.loc[x['LE_ID'] == id]
                y.to_excel(writer, sheet_name=str(id))
                
            else:
                pass
        
        writer.save()
        workbook.close()

如果有任何建议可以加快速度，我将不胜感激。我认为我过度迭代，这导致了这个问题，但我不知道如何解决这个问题。昨天这个函数的版本有点快，但我最终得到了损坏的 XLS 文件——可能是因为我不小心让代码编写器多次访问同一个 xls 文件。

我从这里的社区了解到，列表推导是首选，但我无法找出正确的语法来组织我的函数。我宁愿继续迭代，但消除当前（看似？）冗余迭代。

感谢您的意见

5 月 28 日第一次更新：

数据集有以下字段

LE_ID               object
LEGAL_ENTITY_NAME   object
COUNTRY             object
GL_ACCOUNT                              object
BOOK_AMT                                int64
ADJUSTED_TAX                            float64

5 月 28 日第二次更新：

根据 cmets 中的建议修改代码，完美运行：

def SplitDatasetByCountry(dataset):

for country, country_df in dataset.groupby('COUNTRY'):
      
    ## output
    output_folder = Path('/Users/XXXX/Desktop/Reports')
    output_filename = 'Data_for_' + str(country) + '_.xlsx'
    output = output_folder / output_filename

    ## writer
    writer = pd.ExcelWriter(output)
    workbook = writer.book
    
    ## country logic            
    country_expenses = function_A(country_df)
    country_income = function_B(country_df)
    expense = country_expenses.groupby('LE_ID')['BOOK_AMT'].sum()
    income = country_income..groupby('LE_ID')['BOOK_AMT'].sum()
    expense.to_excel(writer, sheet_name='Country Expense')
    income.to_excel(writer, sheet_name='Country Income')
    
    for le_id, le_id_df in country_df.groupby('LE_ID'):
        le_id_df.to_excel(writer, sheet_name=str(le_id))
                    
    writer.save()
    workbook.close()

【参考方案1】：

您能否在数据框进入您的功能之前分享它的格式？

如果该格式足够好，那么只需根据国家/地区拆分数据框，那么 id 就足够了：

list_of_dfs = [dataset[["COUNTRY"] == country] for country in dataset["COUNTRY"].unique()]

我不明白你想用 id 完成什么，但我至少可以指出，你可以像上面那样在 Series 上使用 .unique() 直接获得唯一性，参见 pandas' doc：

unique_ids = dataset["LE_ID"].unique()

【讨论】：

感谢您向我指出 .unique() 方法——这真的很有帮助。我已经用数据集的关键字段更新了我上面的帖子。你能看看——我相信格式可能“不够好”——正如你所写的——因为你提供的列表理解导致了一个关键错误（关键）。使用您建议的列表理解是否正确，项目“== country”引用“COUNTRY”列中的任何值，即不是对另一个变量的引用？

【参考方案2】：

您需要先按国家分组，然后按 LED_ID 分组并将结果写入文件。

for country, country_df in dataset.groupby('COUNTRY'):
    ## output
    output_folder = Path('/Users/XXXXXX/Desktop/TOOL/Reports')
    output_filename = 'Data__for_' + str(country) + '_.xlsx'
    output = output_folder / output_filename
    ## writer
    with pd.ExcelWriter(output) as writer:
        for led_id, led_df in country_df.groupby("LED_ID"):
            led_df.to_excel(writer, sheet_name=str(led_id))

【讨论】：

感谢 Alexander 的建议，它完美无缺。而我的原始版本需要 40 多分钟才能完成（无疑是由于重复迭代），而您的版本在几分钟内完成。我需要进一步研究 groupby 方法，因为我不太明白发生了什么。感谢您为我指明正确的方向。在上面的原始帖子中上传了修改后的代码以供将来参考。请注意附加函数 A 和 B，它们在数据帧上发挥了一些单独的魔力，并且也被打印到它们自己的工作表上。

【参考方案3】：

有两个术语需要理解。

矢量化和传统循环。

如果你使用 for 循环，甚至嵌套 for 循环。时间复杂度会增加 以 O(n^2) 的顺序。

但如果你使用向量化，你的算法可能会运行得更快。

尝试使用 numpy 求和或点函数来解决问题，而不是 for 循环，这样会更快。

This the example for implementing vectorization techniques