Pandas优化：基于阈值条件计算回溯周期的低效循环改进

嘿，我明白你现在遇到的问题了——这个嵌套循环在数据量变大的时候肯定会慢得让人头疼。先帮你梳理下原代码的逻辑：从后往前遍历每一行，然后逐个往前回溯，只要当前行和回溯行的x差值不超过lb，就检查a是否达到阈值thres，计数直到不满足差值条件或者超出数据范围。

先把你的示例数据整理成清晰的表格，方便对照：

输入输出示例 #

原代码的问题 #

原代码的核心问题是双重循环的时间复杂度太高（O(n²)），数据量一旦超过几千行，运行速度会断崖式下跌。另外，df['x'][i]这种链式索引不仅容易触发SettingWithCopyWarning，访问效率也很低。

下面给你几个优化方案，效率依次提升：

方案1：用Pandas apply简化代码 #

apply是Pandas内部优化过的循环，比手动写Python循环更高效，代码也更简洁：

import pandas as pd
import numpy as np

df = pd.read_csv('input.csv')
lb = 0.25
thres = 5

def count_matching_threshold(row):
    current_idx = row.name
    current_x = row['x']
    count = 0
    # 从当前行往前逐个检查
    for idx in range(current_idx, -1, -1):
        if abs(current_x - df['x'].iloc[idx]) > lb:
            break
        if df['a'].iloc[idx] >= thres:
            count += 1
    return count

# 按行应用函数计算countA
df['countA'] = df.apply(count_matching_threshold, axis=1)

这个方案比原代码简洁不少，而且避免了链式索引的问题，中小数据集下足够用。

方案2：用Numba编译循环（大数据集首选） #

如果你的数据量很大（比如几万行以上），推荐用Numba把循环编译成机器码，速度能提升几十到上百倍：

import pandas as pd
import numpy as np
from numba import jit

df = pd.read_csv('input.csv')
lb = 0.25
thres = 5

# 转为NumPy数组，方便Numba处理
x_arr = df['x'].values
a_arr = df['a'].values
countA_arr = np.zeros(len(x_arr), dtype=np.int64)

# 用Numba编译函数
@jit(nopython=True)
def compute_count(x, a, lb, thres, result):
    n = len(x)
    # 反向遍历每个索引
    for i in range(n-1, -1, -1):
        current_x = x[i]
        count = 0
        j = 0
        while True:
            idx = i - j
            if idx < 0:
                break
            # 差值超过阈值就停止回溯
            if abs(current_x - x[idx]) > lb:
                break
            # 检查a是否达标
            if a[idx] >= thres:
                count += 1
            j += 1
        result[i] = count

# 执行编译后的函数
compute_count(x_arr, a_arr, lb, thres, countA_arr)
df['countA'] = countA_arr

Numba的nopython=True参数会强制函数编译为纯机器码，避免Python对象的开销，是处理大规模循环的利器。

验证结果 #

不管用哪个方案，你都可以用示例数据验证：比如索引5的行（x=0.75，a=9），回溯会统计到自身（a=9）和索引2（a=8）两个达标项，最终countA=2，和预期输出一致。

备注：内容来源于stack exchange，提问作者TrippyB