Pandas条件查找优化：bisect提升效率

IT行业相对于一般传统行业，发展更新速度更快，一旦停止了学习，很快就会被行业所淘汰。所以我们需要踏踏实实的不断学习，精进自己的技术，尤其是初学者。今天golang学习网给大家整理了《Pandas条件索引查找优化：bisect提升效率》，聊聊，我们一起来看看吧！

本文旨在解决Pandas DataFrame中查找满足满足特定条件的最近历史索引的效率问题。针对传统`apply`方法在大数据集上的性能瓶颈，文章详细介绍了如何利用Python内置的`bisect`模块结合字典缓存机制，实现显著的性能提升。通过对比多种方案，`bisect`方法被证明是最优解，为处理此类状态依赖型问题提供了高效且内存友好的解决方案。

1. 引言：理解问题与挑战

在数据分析中，我们经常需要根据当前行的值，从历史数据中查找满足特定条件的记录。一个典型的场景是：给定一个包含lower和upper列以及时间索引DATE的Pandas DataFrame，对于每一行，我们需要找到其之前所有行中，lower值大于或等于当前行upper值的最近一次发生的时间索引。

例如，对于以下DataFrame：

            lower  upper
DATE                    
2020-01-01      7      2
2020-01-02      1      3
2020-01-03      6      4
2020-01-04      1      5
2020-01-05      1      6
2020-01-06      1      7
2020-01-07      1      8
2020-01-08     11      9
2020-01-09      1     10
2020-01-10      1     11

对于2020-01-04这一行，upper值为5。我们需要查找2020-01-04之前的所有行中，lower值大于等于5的最近时间索引。在本例中，2020-01-03的lower值为6 (6 >= 5)，是满足条件的最近索引。

这类问题的一个主要挑战是其固有的“状态依赖性”：当前行的计算结果依赖于之前行的状态，这使得传统的Pandas向量化操作难以直接应用，导致性能成为大数据集上的一个瓶颈。

2. 低效基线方案：DataFrame.apply()

最直观的解决方案是使用DataFrame.apply()方法逐行处理。这种方法虽然易于理解和实现，但其效率极低，尤其是在处理大型DataFrame时。

2.1 方案实现

import pandas as pd
import numpy as np

# 示例DataFrame生成函数
def get_sample_df(rows=10):
    data = {'lower': np.random.default_rng(seed=1).uniform(1,100,rows),
            'upper': np.random.default_rng(seed=2).uniform(1,100,rows)}
    df = pd.DataFrame(data=data).astype(int)
    df['DATE'] = pd.date_range('2020-01-01', periods=rows, freq="min")
    df.set_index('DATE', inplace=True)
    return df

def get_baseline():
    df = get_sample_df()

    def get_most_recent_index(row):
        # 筛选当前行之前的所有行
        previous_indices = df.loc[:row.name - pd.Timedelta(minutes=1)]  
        # 在之前行中找到满足条件的行，并返回最近的索引
        recent_index = previous_indices[previous_indices['lower'] >= row['upper']].index.max()
        return recent_index

    df['prev'] = df.apply(get_most_recent_index, axis=1) 
    return df

# 运行示例
df_baseline = get_baseline()
print(df_baseline)

2.2 性能分析

上述apply方法效率低下的主要原因在于：

1. 逐行迭代：apply(axis=1)本质上是Python级别的循环，无法利用Pandas底层的C优化。
2. 重复切片：在每次迭代中，df.loc[:row.name - pd.Timedelta(minutes=1)]都会对DataFrame进行切片操作，这会创建新的DataFrame视图或副本，开销巨大。
3. 重复筛选：previous_indices[previous_indices['lower'] >= row['upper']]在每次迭代中都会重新执行条件筛选。

对于包含10万行数据的DataFrame，此方法的执行时间可能长达数分钟，甚至更久。

3. 高效解决方案：利用二分查找 (bisect)

为了显著提升性能，我们需要避免重复的DataFrame切片和筛选操作，并利用更高效的数据结构和算法。Python的内置bisect模块提供二分查找功能，结合一个字典来缓存已见过的lower值及其最近日期，可以实现高效查找。

3.1 bisect模块简介

bisect模块实现了一个二分查找算法，用于在有序序列中查找插入点，以保持序列的有序性。bisect_left(a, x)函数返回在有序序列a中插入x后，x仍然保持有序的左侧插入点索引。这意味着所有a[i]，其中i < bisect_left(a, x)，都小于x。

3.2 方案实现

核心思想是：

1. 维护一个已排序的唯一lower值列表 (uniq_lower)，用于二分查找。
2. 维护一个字典 (last_seen)，存储每个lower值最近一次出现的日期。
3. 对于每一行：
   • 使用bisect_left在uniq_lower中找到所有大于或等于当前行upper值的lower值的起始位置。
   • 遍历这些符合条件的lower值，从last_seen字典中获取它们对应的最近日期。
   • 选择这些日期中的最大值（即最近的日期）作为结果。
   • 将当前行的lower值和日期更新到last_seen字典中。

from bisect import bisect_left

def get_bisect():
    df = get_sample_df() # 使用相同的示例数据生成函数

    def get_prev_bs(lower_series, upper_series, date_index):
        # 存储所有出现过的唯一lower值，并保持排序
        uniq_lower = sorted(list(set(lower_series)))
        # 存储每个lower值最近一次出现的日期
        last_seen = {}

        results = []
        for l, u, d in zip(lower_series, upper_series, date_index):
            # 使用二分查找找到在uniq_lower中，第一个大于或等于u的元素的索引
            # 这意味着uniq_lower[idx:]包含了所有 >= u 的lower值
            idx = bisect_left(uniq_lower, u)

            max_date = None
            # 遍历所有符合条件的lower值
            for lv in uniq_lower[idx:]:
                if lv in last_seen:
                    # 如果该lower值之前出现过
                    if max_date is None:
                        max_date = last_seen[lv]
                    elif last_seen[lv] > max_date:
                        # 更新为更近的日期
                        max_date = last_seen[lv]
            results.append(max_date)
            # 更新当前lower值最近一次出现的日期
            last_seen[l] = d
        return results

    df["prev"] = list(get_prev_bs(df["lower"], df["upper"], df.index))
    return df

# 运行示例
df_bisect = get_bisect()
print(df_bisect)

3.3 结果验证

使用原始问题中的示例数据进行验证：

import pandas as pd
from bisect import bisect_left

data = {'lower': [7, 1, 6, 1, 1, 1, 1, 11, 1, 1],
        'upper': [2, 3, 4, 5, 6, 7, 8, 9, 10, 11]}
df = pd.DataFrame(data=data)
df['DATE'] = pd.date_range('2020-01-01', periods=len(data['lower']))
df.set_index('DATE', inplace=True)

def get_prev_bs_verify(lower_series, upper_series, date_index):
    uniq_lower = sorted(list(set(lower_series)))
    last_seen = {}
    results = []
    for l, u, d in zip(lower_series, upper_series, date_index):
        idx = bisect_left(uniq_lower, u)
        max_date = None
        for lv in uniq_lower[idx:]:
            if lv in last_seen:
                if max_date is None:
                    max_date = last_seen[lv]
                elif last_seen[lv] > max_date:
                    max_date = last_seen[lv]
        results.append(max_date)
        last_seen[l] = d
    return results

df["prev_new"] = list(get_prev_bs_verify(df["lower"], df["upper"], df.index))
print(df)

输出：

            lower  upper   prev_new
DATE                             
2020-01-01      7      2        NaT
2020-01-02      1      3 2020-01-01
2020-01-03      6      4 2020-01-01
2020-01-04      1      5 2020-01-03
2020-01-05      1      6 2020-01-03
2020-01-06      1      7 2020-01-01
2020-01-07      1      8        NaT
2020-01-08     11      9        NaT
2020-01-09      1     10 2020-01-08
2020-01-10      1     11 2020-01-08

结果与预期一致。

4. 其他尝试与性能对比

除了上述两种方法，还有其他一些尝试，例如使用pyjanitor库或基于纯Python列表的enumerate循环。然而，这些方法在性能或内存效率上存在局限性。

4.1 pyjanitor方案（内存限制）

pyjanitor库提供了conditional_join等功能，旨在进行条件连接。虽然在某些场景下能提供向量化优势，但对于本例中涉及的复杂条件和大量数据，它可能导致巨大的中间数据结构，从而引发内存分配错误。

4.2 enumerate方案（效率低下）

此方案将DataFrame转换为Python列表，然后使用嵌套循环进行迭代和条件判断。虽然避免了DataFrame切片，但其核心仍是Python级别的循环，并且内部的any()和reversed()操作在每次迭代中都会重新遍历列表切片，导致效率低下。

4.3 性能测试结果

对包含10万行数据的DataFrame进行性能测试，结果如下：

从结果可以看出，bisect方案以压倒性的优势胜出，其速度比baseline和enumerate方案快了近60倍。pyjanitor方案则因内存限制未能完成测试。

5. 注意事项与最佳实践

1. 理解问题本质：当问题涉及“基于历史状态的逐行计算”时，直接的Pandas向量化通常难以实现。此时，需要转向更底层的Python循环，但必须辅以高效的算法和数据结构。
2. 利用内置模块：Python标准库提供了许多优化工具，如bisect、heapq等，它们针对特定任务进行了高度优化。在面临性能瓶颈时，考虑这些内置工具往往能带来惊喜。
3. 时间复杂度分析：
   • baseline方案：对于N行数据，每行都进行DataFrame切片和筛选，大致为O(N^2)甚至更高。
   • bisect方案：初始化uniq_lower为O(N log N)（排序）。主循环中，每次迭代bisect_left是O(log M)（M是uniq_lower的长度），内部遍历uniq_lower[idx:]最坏情况是O(M)。因此，整体复杂度约为O(N log N + N * M)。在lower值种类不多的情况下，M远小于N，此方案非常高效。
4. 内存管理：对于大数据集，避免创建大型中间数据结构至关重要。bisect方案通过维护一个last_seen字典和uniq_lower列表，其内存开销相对稳定且可控。

6. 总结

在Pandas中处理依赖于历史状态的条件查找问题时，直接使用DataFrame.apply()虽然简单但效率低下。通过将问题分解，并利用Python内置的bisect模块结合字典缓存机制，可以构建一个高度优化的解决方案。这种方法不仅显著提升了计算速度，还有效地管理了内存开销，使其成为处理大规模数据集此类问题的最佳实践。对于需要从历史数据中快速检索满足特定条件的记录的场景，bisect方案提供了一个强大且高效的工具。

好了，本文到此结束，带大家了解了《Pandas条件查找优化：bisect提升效率》，希望本文对你有所帮助！关注golang学习网公众号，给大家分享更多文章知识！