R data.table场景下按子集求和更新的高效实现方案咨询

优化思路 #

你要实现的是集合上的莫比乌斯反演运算，原代码慢的核心问题是逐行R级循环+行级向量运算，完全没有利用data.table的向量化操作优势，我们可以通过位掩码编码集合+批量超集判断+分组聚合的思路将运算速度提升1~2个数量级。

实现代码 #

library(data.table)

# 测试数据
sets <- data.table(
    id.s = c(1,2,2,3,3,3,4,4,4,4),
    id.e = c(3,3,4,2,3,4,1,2,3,4))
v <- data.table(id.s = 1:4, value = c(1/10,2/10,3/10,4/10))

# 步骤1：给每个元素分配位索引，生成集合的位掩码与基数
elem_map <- data.table(id.e = unique(sets$id.e))[, bit_idx := .I - 1]
sets <- sets[elem_map, on = "id.e"]
set_info <- sets[, .(
  mask = sum(bitwShiftL(1, bit_idx)),
  cardinality = .N
), by = id.s][v, on = "id.s"]

# 步骤2：生成集合对，过滤超集关系，聚合计算结果
pair_dt <- CJ(id.s.a = set_info$id.s, id.s.b = set_info$id.s)
pair_dt <- pair_dt[set_info, on = c("id.s.a" = "id.s")][
  set_info, on = c("id.s.b" = "id.s")
]
pair_dt <- pair_dt[bitwAnd(mask, i.mask) == mask]
res <- pair_dt[, .(
  value2 = sum((-1)^(i.cardinality - cardinality) * i.value)
), by = .(id.s = id.s.a)]

print(res)
#    id.s value2
# 1:    1   -0.2
# 2:    2    0.3
# 3:    3   -0.1
# 4:    4    0.4

适配大尺寸数据的补充方案 #

• 若元素个数≤64：替换bitwShiftL、bitwAnd为bit64包的64位整数运算函数即可
• 若元素个数>64：可以用bit包的bit向量存储集合，或者将集合转成Matrix包的稀疏0-1矩阵，通过矩阵运算批量判断包含关系，依然远快于原循环方案
• 若集合数超过1万：可以按基数从大到小排序后用动态规划计算，时间复杂度可优化至O(集合数*平均集合大小)

内容的提问来源于stack exchange，提问作者Vaclav Kratochvíl