在R中基于索引计算两矩阵间加权欧氏距离的高效方法

问题说明 #

有两个存储产品分组信息的数据集：

• 小数据集df_with_group_values：每组对应一行，包含1列分组索引和若干数值变量
• 大数据集df_with_data：每组包含多行数据，包含产品ID、分组索引和相同的数值变量

需要针对每个产品，计算其所在分组在小数据集中的对应行与该产品行之间的加权欧氏距离。现有for循环实现在分组数量增加时运行速度变慢，需更高效优雅的方案。

示例数据 #

# 修正原代码的rnorm参数顺序问题
df_with_group_values = data.frame(
  group= 1:6, 
  matrix(rnorm(30, mean=100), nrow=6))

df_with_data = data.frame(
  product_id = 1:1000,
  group = sample(6, 1000, replace=T),
  matrix(rnorm(5000, mean=100), nrow=1000))

# 定义权重向量
weights = rnorm(5)

高效实现方案 #

方法1：用dplyr实现（优雅易读） #

核心思路是通过分组列连接两个数据集，利用向量化运算批量计算距离，避免循环：

library(dplyr)

# 给数值列统一命名，方便后续操作
names(df_with_group_values)[-1] <- paste0("var_", 1:5)
names(df_with_data)[-c(1,2)] <- paste0("var_", 1:5)

df_distancias <- df_with_data %>%
  # 按group关联分组基准数据
  left_join(df_with_group_values, by = "group", suffix = c("", "_group")) %>%
  # 计算每个变量的加权平方差
  mutate(across(
    .cols = starts_with("var_"),
    .fns = ~ weights[match(cur_column(), paste0("var_", 1:5))] * (.x - get(paste0(cur_column(), "_group")))^2,
    .names = "weighted_sq_{.col}"
  )) %>%
  # 计算加权欧氏距离（若需加权平均距离，将rowSums替换为rowMeans）
  mutate(distance = sqrt(rowSums(select(., starts_with("weighted_sq_"))))) %>%
  # 保留目标列并排序
  select(product_id, distance) %>%
  arrange(product_id)

方法2：用data.table实现（性能最优） #

针对超大规模数据集，data.table的内存效率和运算速度远高于普通循环和dplyr：

library(data.table)

# 转换为data.table格式
setDT(df_with_group_values)
setDT(df_with_data)

# 给数值列统一命名
setnames(df_with_group_values, paste0("var_", 1:5), skip = 1)
setnames(df_with_data, paste0("var_", 1:5), skip = 2)

# 关联数据+计算距离
df_distancias <- df_with_data[df_with_group_values, on = "group", allow.cartesian = TRUE][
  , 
  # 批量生成加权平方差列
  (paste0("weighted_sq_", 1:5)) := lapply(1:5, function(i) weights[i] * (get(paste0("var_", i)) - get(paste0("var_", i), envir = parent.frame(2)))^2)
][
  ,
  # 计算距离
  distance := sqrt(rowSums(.SD)), .SDcols = paste0("weighted_sq_", 1:5)
][
  # 保留目标列并排序
  order(product_id), .(product_id, distance)
]

原循环方案的效率问题 #

• 循环中频繁筛选数据、创建临时矩阵，产生大量不必要的内存开销
• 用rbind逐步拼接数据框，每次都会复制整个对象，时间复杂度呈O(n²)增长
• 没有利用R的向量化运算优势，原生循环本身效率低下

内容的提问来源于stack exchange，提问作者Álvaro Méndez Civieta