如何移除R data.table中基于子集最小mass填充NA的循环？

我有一个data.table，每一行对应一个带有mass赋值的集合。有时mass的值为NA，我希望用该集合所有子集的非NAmass的最小值来填充这些NA。

集合由set列（整数位标记）编码，imprint列是集合的位图表示（用于可视化检查），N代表集合的基数。

示例数据：

dt <- data.table(id.fe = 1:15, set = c(1,2,4,8,3,5,9,6,10,12,7,11,13,14,15), imprint = c('0001','0010','0100','1000','0011','0101','1001','0110','1010','1100','0111','1011','1101','1110','1111'), N = c(1,1,1,1,2,2,2,2,2,2,3,3,3,3,4),  mass = c(0.4,1,1,0,0.3,NA,NA,NA,NA,0,NA,NA,NA,NA,NA))

数据预览：

id.fe set imprint N mass
 1:     1   1    0001 1  0.4
 2:     2   2    0010 1  1.0
 3:     3   4    0100 1  1.0
 4:     4   8    1000 1  0.0
 5:     5   3    0011 2  0.3
 6:     6   5    0101 2   NA
 7:     7   9    1001 2   NA
 8:     8   6    0110 2   NA
 9:     9  10    1010 2   NA
10:    10  12    1100 2  0.0
11:    11   7    0111 3   NA
12:    12  11    1011 3   NA
13:    13  13    1101 3   NA
14:    14  14    1110 3   NA
15:    15  15    1111 4   NA

比如id.fe==6的二元集合，mass为NA，需要用它的子集（id.fe==1和id.fe==3的集合）的最小mass（0.4）来替换，结果存到新列newMass中。

目前我用for循环实现这个逻辑：

for(id in dt[is.na(mass),id.fe]) {
  idSet <- dt[id.fe == id, set]
  idN <- dt[id.fe == id, N]
  dt[id.fe == id, newMass := min(dt[N < idN & !is.na(mass) & bitwAnd(set, idSet) == set, ]$mass)]
}

运行后结果：

id.fe set imprint N mass newMass
 1:     1   1    0001 1  0.4      NA
 2:     2   2    0010 1  1.0      NA
 3:     3   4    0100 1  1.0      NA
 4:     4   8    1000 1  0.0      NA
 5:     5   3    0011 2  0.3      NA
 6:     6   5    0101 2   NA     0.4
 7:     7   9    1001 2   NA     0.0
 8:     8   6    0110 2   NA     1.0
 9:     9  10    1010 2   NA     0.0
10:    10  12    1100 2  0.0      NA
11:    11   7    0111 3   NA     0.3
12:    12  11    1011 3   NA     0.0
13:    13  13    1101 3   NA     0.0
14:    14  14    1110 3   NA     0.0
15:    15  15    1111 4   NA     0.0

请问有没有办法移除这个循环？

可以利用data.table的向量化连接和分组聚合特性实现无循环操作，以下是两种高效实现方式：

方法1：直接分组计算（代码简洁） #

先筛选出非NA的mass记录作为候选子集，再对每个需要填充的行，直接匹配符合条件的子集并计算最小值：

# 提取非NA的mass记录作为候选子集
non_na_mass <- dt[!is.na(mass)]

# 对每个NA行，计算其所有真子集的最小mass
dt[is.na(mass), newMass := non_na_mass[bitwAnd(set, .BY$set) == set & N < .BY$N, min(mass)], by = .(set, N)]

方法2：预先生成子集关系（适合大数据场景） #

如果数据集较大，预先生成所有合法的子集-父集配对，再做聚合计算，能进一步提升效率：

# 生成所有子集与父集的配对，仅保留父集基数更大的真子集关系
subset_pairs <- dt[!is.na(mass), .(subset_set = set, subset_mass = mass)][
  dt, on = .(subset_set = set), allow.cartesian = TRUE
][bitwAnd(subset_set, set) == subset_set & N > i.N]

# 分组计算每个父集的最小子集mass
min_subset_mass <- subset_pairs[, .(newMass = min(subset_mass)), by = .(set, N)]

# 将结果合并回原表
dt[min_subset_mass, on = .(set, N), newMass := i.newMass]

逻辑说明 #

• 两种方法都完全避免了循环，利用data.table的向量化操作提升计算效率
• bitwAnd(subset_set, set) == subset_set确保子集的位标记完全包含在父集内，符合子集定义
• 过滤N < .BY$N（或N > i.N）保证只考虑基数更小的真子集
• 最终通过分组聚合得到每个需要填充的集合的最小子集mass

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