如何优化Box-Cox变换的Python函数以提升大数据量运行速度？

问题背景 #

我编写了一个计算给定序列lambda系数的函数，但处理大尺寸输入时运行耗时过长，求提速方法。原始代码如下：

from scipy.stats import norm, pearsonr
import numpy as np  # 原代码漏导入，需补上

def get_lambda_coef(series):
    x=[series[i] for i in range(len(series))]
    for i in range(len(x)-1):
        for j in range(len(x)-1):
            if x[j]>=x[j+1]:
                z=x[j]
                x[j]=x[j+1]
                x[j+1]=z
    i=[j for j in range(1,len(x)+1)]
    f=[(i[j]-0.375)/(len(x)+0.25) for j in range(len(x))]
    u=[norm.ppf(f[i]) for i in range(len(x))]
        
    lambda_coef=0
    width=3
    step=width/6
    k=lambda_coef-width
    iteration=1
    while iteration<=15:
        r_vector=[]
        lambda_vect=[]
        while k<=lambda_coef+width:
            if k==0:
                y=[np.log(i) for i in x]
            else:
                y=[(i**k-1)/k for i in x]
            r_vector.append(pearsonr(y, u)[0])
            k+=step
        k=lambda_coef-width
        while k<=lambda_coef+width:
            lambda_vect.append(k)
            k+=step
        lambda_coef=lambda_vect[r_vector.index(max(r_vector))]
        width/=2
        step/=3
        k=lambda_coef-width
        iteration+=1
    normalized = [(x**lambda_coef - 1)/lambda_coef for x in series]
    return (normalized, lambda_coef)

具体提速技巧 #

• 替换低效排序算法：原代码用双重循环实现冒泡排序，时间复杂度为O(n²)，对大数据量极不友好。直接改用Python内置的list.sort()（时间复杂度O(n log n)）或numpy的np.sort()，速度提升显著：

  # 替换原手动排序代码
  x = sorted(series)  # 或转numpy数组后排序：x = np.sort(np.array(series))
  

• 用numpy向量化操作替代列表推导式：Python列表推导是纯Python循环，numpy的向量化操作基于C实现，速度快得多。将所有列表操作转为numpy数组操作：

  # 替换x、i、f、u的生成
  x_np = np.sort(np.array(series))
  n = len(x_np)
  i_np = np.arange(1, n+1)
  f_np = (i_np - 0.375) / (n + 0.25)
  u_np = norm.ppf(f_np)
  

• 避免重复计算lambda_vect：原代码用两次while循环分别生成r_vector和lambda_vect，可一次性生成lambda_vect数组，再批量计算r_vector：

  # 替换迭代中的循环
  lambda_vect = np.arange(lambda_coef - width, lambda_coef + width + step, step)
  r_vector = []
  for k in lambda_vect:
      if np.isclose(k, 0):
          y_np = np.log(x_np)
      else:
          y_np = (x_np ** k - 1) / k
      # 用numpy计算皮尔逊相关系数，比scipy的pearsonr略快
      corr = np.corrcoef(y_np, u_np)[0, 1]
      r_vector.append(corr)
  

• 优化迭代终止条件：原代码固定15次迭代，可改为当lambda_coef的变化量小于阈值（如1e-6）时提前终止，减少不必要的计算：

  # 替换原迭代逻辑
  lambda_coef = 0.0
  width = 3.0
  step = width / 6
  tol = 1e-6
  prev_lambda = lambda_coef
  iteration = 0
  max_iter = 15
  while iteration < max_iter:
      lambda_vect = np.arange(lambda_coef - width, lambda_coef + width + step, step)
      r_vector = []
      for k in lambda_vect:
          if np.isclose(k, 0):
              y_np = np.log(x_np)
          else:
              y_np = (x_np ** k - 1) / k
          corr = np.corrcoef(y_np, u_np)[0, 1]
          r_vector.append(corr)
      best_idx = np.argmax(r_vector)
      lambda_coef = lambda_vect[best_idx]
      # 检查是否收敛
      if abs(lambda_coef - prev_lambda) < tol:
          break
      prev_lambda = lambda_coef
      width /= 2
      step /= 3
      iteration += 1
  

• 减少函数调用开销：原代码每次循环都调用pearsonr，改用numpy的np.corrcoef可减少函数调用的额外开销；另外，将循环内的重复计算（如len(x)）提前赋值给变量，避免重复计算。

• 最终向量化的归一化计算：最后生成normalized时，用numpy数组操作替代列表推导：

  if np.isclose(lambda_coef, 0):
      normalized = np.log(np.array(series))
  else:
      normalized = (np.array(series) ** lambda_coef - 1) / lambda_coef
  

优化后的完整代码示例 #

import numpy as np
from scipy.stats import norm

def get_lambda_coef_optimized(series):
    # 转为numpy数组并排序
    x_np = np.sort(np.array(series))
    n = len(x_np)
    if n == 0:
        raise ValueError("Series cannot be empty")
    
    # 计算f和u的向量化操作
    i_np = np.arange(1, n+1)
    f_np = (i_np - 0.375) / (n + 0.25)
    u_np = norm.ppf(f_np)
    
    lambda_coef = 0.0
    width = 3.0
    step = width / 6
    tol = 1e-6
    prev_lambda = lambda_coef
    max_iter = 15
    iteration = 0
    
    while iteration < max_iter:
        # 生成lambda候选值数组
        lambda_vect = np.arange(lambda_coef - width, lambda_coef + width + step, step)
        r_vector = []
        
        for k in lambda_vect:
            if np.isclose(k, 0):
                y_np = np.log(x_np)
            else:
                y_np = (x_np ** k - 1) / k
            # 计算皮尔逊相关系数
            corr = np.corrcoef(y_np, u_np)[0, 1]
            r_vector.append(corr)
        
        # 找到最大相关系数对应的lambda
        best_idx = np.argmax(r_vector)
        lambda_coef = lambda_vect[best_idx]
        
        # 收敛检查
        if abs(lambda_coef - prev_lambda) < tol:
            break
        
        prev_lambda = lambda_coef
        width /= 2
        step /= 3
        iteration += 1
    
    # 生成归一化序列
    series_np = np.array(series)
    if np.isclose(lambda_coef, 0):
        normalized = np.log(series_np)
    else:
        normalized = (series_np ** lambda_coef - 1) / lambda_coef
    
    return (normalized, lambda_coef)

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