如何用data.table填充缺失值？基于曲线拟合的解决方案

我有两个data.table对象，其中y列存在缺失值，非缺失值构成一条连续曲线，需要将缺失值贴合该曲线进行填充。以下是示例数据：

示例数据1 #

temp_dt_1 = structure(list(x = c(33, 34, 35, 36, 37, 38, 39, 40, 45, 50, 
51.5, 52.5, 53, 54, 55, 55.5, 56, 56.5, 57, 57.5, 58, 59, 59.5, 
60, 60.5, 61, 61.5, 62, 62.5, 63, 63.5, 64, 65, 66, 66.5, 67, 
67.5), y = c(NA, NA, NA, NA, NA, NA, NA, NA, NA, NA, NA, NA, 
NA, NA, NA, NA, NA, 2.33, NA, NA, 1.6, NA, NA, 0.91, NA, 0.665, 
NA, 0.465, 0.39, 0.325, 0.265, 0.24, 0.135, 0.1, 0.095, NA, 0.105
)), row.names = c(NA, -37L), class = c("data.table", "data.frame"
))

数据预览：

x     y
 1: 33.0    NA
 2: 34.0    NA
 3: 35.0    NA
 4: 36.0    NA
 5: 37.0    NA
 6: 38.0    NA
 7: 39.0    NA
 8: 40.0    NA
 9: 45.0    NA
10: 50.0    NA
11: 51.5    NA
12: 52.5    NA
13: 53.0    NA
14: 54.0    NA
15: 55.0    NA
16: 55.5    NA
17: 56.0    NA
18: 56.5 2.330
19: 57.0    NA
20: 57.5    NA
21: 58.0 1.600
22: 59.0    NA
23: 59.5    NA
24: 60.0 0.910
25: 60.5    NA
26: 61.0 0.665
27: 61.5    NA
28: 62.0 0.465
29: 62.5 0.390
30: 63.0 0.325
31: 63.5 0.265
32: 64.0 0.240
33: 65.0 0.135
34: 66.0 0.100
35: 66.5 0.095
36: 67.0    NA
37: 67.5 0.105

示例数据2（首尾均有缺失） #

temp_dt_2 = structure(list(x = c(25, 27.5, 30, 32.5, 35, 40, 43, 45, 47, 
50, 52.5, 55, 57.5, 60, 65, 85, 90), y = c(NA, NA, 1.1, NA, 
1.765, 2.64, 3.42, 3.975, NA, 5.625, 6.95, 7.825, 8.95, 10.375, 
13.725, NA, NA)), row.names = c(NA, -17L), class = c("data.table", 
"data.frame"))

预期效果 #

已知非缺失点构成基础曲线，填充后的缺失值需贴合该曲线趋势（如示例1中预期的连续曲线所示）。

解决方案（仅用data.table） #

以下是高性能的填充函数，完全基于data.table内置函数实现，避免循环，适合大数据量场景：

library(data.table)

fill_curve_na <- function(dt) {
  # 提取非缺失的参考点并按x排序
  ref_points <- dt[!is.na(y), .(x, y)]
  setorder(ref_points, x)
  
  # 为每个x定位前后相邻的参考点索引
  dt[, `:=`(
    prev_x_idx = findInterval(x - 1e-9, ref_points$x),
    next_x_idx = findInterval(x, ref_points$x)
  )]
  
  # 中间缺失值：线性插值保证曲线连续
  dt[is.na(y) & prev_x_idx > 0 & next_x_idx <= nrow(ref_points), 
     y := ref_points$y[prev_x_idx] + (ref_points$y[next_x_idx] - ref_points$y[prev_x_idx]) * 
       (x - ref_points$x[prev_x_idx]) / (ref_points$x[next_x_idx] - ref_points$x[prev_x_idx])]
  
  # 处理开头缺失（x小于第一个参考点）：用前两个参考点的斜率外推
  if (nrow(ref_points) >= 2) {
    start_slope <- (ref_points$y[2] - ref_points$y[1]) / (ref_points$x[2] - ref_points$x[1])
    dt[is.na(y) & x < ref_points$x[1], 
       y := ref_points$y[1] + start_slope * (x - ref_points$x[1])]
  } else if (nrow(ref_points) == 1) {
    # 仅一个已知点时，所有缺失值填充为该点的y值
    dt[is.na(y), y := ref_points$y[1]]
  }
  
  # 处理结尾缺失（x大于最后一个参考点）：用最后两个参考点的斜率外推
  if (nrow(ref_points) >= 2) {
    end_slope <- (ref_points$y[nrow(ref_points)] - ref_points$y[nrow(ref_points)-1]) / 
      (ref_points$x[nrow(ref_points)] - ref_points$x[nrow(ref_points)-1])
    dt[is.na(y) & x > ref_points$x[nrow(ref_points)], 
       y := ref_points$y[nrow(ref_points)] + end_slope * (x - ref_points$x[nrow(ref_points)])]
  }
  
  # 清理临时计算列
  dt[, c("prev_x_idx", "next_x_idx") := NULL]
  return(dt)
}

测试代码 #

# 填充示例1
filled_dt1 <- fill_curve_na(copy(temp_dt_1))
print(filled_dt1)

# 填充示例2
filled_dt2 <- fill_curve_na(copy(temp_dt_2))
print(filled_dt2)

方案说明 #

• 高性能：使用findInterval快速定位参考点，避免逐行循环，处理大数据集效率更高
• 贴合趋势：中间缺失值用线性插值保证曲线连续，首尾缺失值用相邻参考点的斜率外推，贴合原有曲线趋势
• 鲁棒性：处理了仅存在单个已知点的极端情况

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