如何用快速傅里叶变换(FFT)加速数据框嵌套循环计算？

核心逻辑拆解 #

你的my_function是各维度RBF核的乘积，最终要计算的是df1每行与df2所有行的核函数值之和。这个计算可以拆解为：对每个维度单独计算df1每行与df2所有行的RBF核值之和，再将各维度的结果对应相乘。

而单个维度的求和操作，本质是卷积运算——因为RBF核是对称的（(y-x)²=(x-y)²），sum_{y∈df2} K(y-x)等价于df2的"信号"（每个y对应权重1）与核函数K的卷积，在x处的取值。FFT可以将卷积的时间复杂度从O(n²)降到O(n log n)。

具体实现步骤 #

1. 单个维度的FFT卷积计算

先实现一个函数，对单个维度计算df1每行对应的核值之和：

compute_dim_sum <- function(col1, col2) {
  # 获取当前维度的所有取值范围，计算最大可能的差值
  all_vals <- c(col1, col2)
  min_val <- min(all_vals)
  max_val <- max(all_vals)
  max_diff <- max_val - min_val
  
  # 构造RBF核：覆盖所有可能的差值（从 -max_diff 到 max_diff）
  t_vals <- seq(-max_diff, max_diff)
  kernel <- exp(log(0.5) * t_vals^2)
  
  # 将列值转换为相对于min_val的偏移（从0开始的索引）
  col1_offset <- col1 - min_val
  col2_offset <- col2 - min_val
  
  # 构造df2的信号向量：统计每个取值出现的次数
  signal_length <- max_val - min_val + 1
  signal <- numeric(signal_length)
  for (y_offset in col2_offset) {
    signal[y_offset + 1] <- signal[y_offset + 1] + 1
  }
  
  # 补零避免循环卷积，长度为核长度 + 信号长度 - 1
  pad_length <- length(kernel) + signal_length - 1
  signal_padded <- c(signal, numeric(pad_length - signal_length))
  kernel_padded <- c(kernel, numeric(pad_length - length(kernel)))
  
  # FFT计算卷积并逆变换
  fft_signal <- fft(signal_padded)
  fft_kernel <- fft(kernel_padded)
  fft_conv <- fft_signal * fft_kernel
  conv_result <- Re(fft(fft_conv, inverse = TRUE)) / pad_length
  
  # 提取df1每行对应的结果
  conv_result[col1_offset + max_diff + 1]
}

2. 多维度结果合并

对每个维度调用上述函数，再将各维度的结果对应相乘：

df1 <- data.frame(a=c(1,1,1,1,1,1,1,1), b=c(1,2,3,1,2,3,1,2), c=c(0,0,1,0,0,1,0,0))
df2 <- data.frame(a=c(1,2,1,2,1,2,1,2), b=c(1,2,3,1,2,3,1,2), c=c(0,0,1,0,0,1,0,0))

n <- nrow(df1)
p <- ncol(df1)

# 计算每个维度的求和结果
dim_results <- lapply(seq_len(p), function(d) {
  compute_dim_sum(df1[[d]], df2[[d]])
})

# 合并各维度结果：对应元素相乘
final_result <- Reduce(`*`, dim_results)

# 输出结果（与原循环结果一致，浮点误差范围内）
final_result

复杂度说明 #

每个维度的FFT计算复杂度为O(M log M)，其中M是补零后的向量长度（远小于n²）。p个维度的总复杂度为O(p*M log M)，当n较大时，这个复杂度远低于原O(n²)的嵌套循环，优化效果明显。

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