如何用data.table筛选指定列仅含一个NA且TOTAL非NA的行并补全NA？

现有如下data.table数据：

library(data.table)
cols <- c("A", "B", "C")
dt <- data.table(A = c(NA,10,20,15),
                 B = c(30,10,5,15),
                 C = c(40,10,NA,15),
                 TOTAL = c(75,30,30,45))

数据输出为：

A  B  C TOTAL
1: NA 30 40    75
2: 10 10 10    30
3: 20  5 NA    30
4: 15 15 15    45

需求：

• 筛选出满足两个条件的行：TOTAL列值非NA，且cols指定的A、B、C列中恰好有一个NA值（即示例中的第1、3行）
• 对筛选出的行，用TOTAL减去该行其他非NA列的和，补全NA值，最终得到结果：

A  B  C TOTAL
1:  5 30 40    75
2: 20  5  5    30

已尝试代码：

dt[!is.na(TOTAL), rowSums(is.na(.SD)) == 1, .SDcols = cols]

需要完整实现筛选与补全功能。

可以通过data.table的链式操作一步完成筛选和补全，代码如下：

library(data.table)
cols <- c("A", "B", "C")
dt <- data.table(A = c(NA,10,20,15),
                 B = c(30,10,5,15),
                 C = c(40,10,NA,15),
                 TOTAL = c(75,30,30,45))

# 筛选并补全NA值
result <- dt[!is.na(TOTAL) & rowSums(is.na(.SD)) == 1, .SDcols = cols][
  , (cols) := lapply(cols, function(x) {
    ifelse(is.na(get(x)), TOTAL - rowSums(.SD, na.rm = TRUE), get(x))
  }), .SDcols = cols
]

print(result)

代码解释 #

1. 筛选行：
   • !is.na(TOTAL) 确保TOTAL列无缺失值
   • rowSums(is.na(.SD)) == 1 统计cols指定列中NA的数量，只保留恰好有1个NA的行，.SDcols = cols限定计算范围为A、B、C列
2. 补全NA值：
   • 对cols中的每一列，用ifelse判断是否为NA：
     • 如果是NA，用TOTAL减去该行其他非NA列的和（rowSums(.SD, na.rm = TRUE)中na.rm = TRUE忽略NA值求和）
     • 如果不是NA，保留原数值
   • 使用:=操作符直接在筛选后的结果中修改列值，符合data.table的高效操作风格

运行后输出结果：

A  B  C TOTAL
1:  5 30 40    75
2: 20  5  5    30

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