在 data.table 中查找最接近的值

Question

我正在尝试为每个处理过的观察找到最接近的值。数据如下（1.2M obs的部分数据）：

> dta
     id treatment      score
 1:   5         0 0.02381024
 2:  10         0 0.05428605
 3:  22         0 0.02118124
 4:  27         0 0.01495214
 5:  45         0 0.01877916
 6:  50         0 0.02120360
 7:  58         0 0.02207263
 8:  60         0 0.02807019
 9:  61         0 0.05432927
10:  65         1 0.59612077
11:  68         0 0.02482168
12:  72         1 0.14582400
13:  73         0 0.02371670
14:  77         0 0.02608826
15:  87         0 0.06852409
16:  88         0 0.07473471
17:  94         0 0.07160314
18:  97         0 0.02040747
19: 104         1 0.09878789
20: 108         0 0.02421807

对于每个处理过的观察结果（即处理 = 1），我希望获得一个未处理的观察值（即处理 = 0），其得分最接近，并将所选观察值标记为不可用于其他处理过的观察值匹配。

例如，第一个处理的观察（第 10 行）将匹配到 id = 88（第 16 行），第 12 行到第 17 行，依此类推。目前我正在运行流动循环：

smpl_treated = dta[treatment == 1]
smpl_untreated = dta[treatment == 0]

n_tmp = nrow(smpl_treated)
matched_id = matrix(0, n_tmp, 1)

smpl_tmp = smpl_untreated

for (i in 1:nrow(smpl_treated)) {

  x = smpl_treated[i]$score

  setkey(smpl_tmp, score)
  tmp = smpl_tmp[J(x), roll = "nearest"]
  matched_id[i] = tmp[[1]]
  smpl_tmp = smpl_tmp[id != tmp[[1]]]

}

matched_smpl = smpl_untreated[id %in% matched_id]

> matched_smpl
   id treatment      score
1: 87         0 0.06852409
2: 94         0 0.07160314
3: 88         0 0.07473471

有什么建议可以在 data.table 中实现或加快循环速度？使用原始的 1.2M obs 循环需要 2 多个小时。提前感谢您的帮助！

Answer 1

这详细说明了the already accepted answer of denis 使用data.table 语法的实际 可能性，例如，在加入时使用on 参数而不是setkey()。

# determine the minimum number of treated and untreated cases
n <- min(dta[treatment == 0L, .N], dta[treatment == 1L, .N])
# order by descending score 
mdt <- dta[order(-score)][
  # and pick the ids of the top n treated and untreated cases
  # so that the highest untreated score match the highest treated score,
  # the 2nd highest untreated the 2nd highest treated and so forth
  , .(id0 = head(.SD[treatment == 0L, id], n), id1 = head(.SD[treatment == 1L, id], n))]
mdt

   id0 id1
1:  88  65
2:  94  72
3:  87 104

# join the ids two times to show the data of the treated and untreated cases
dta[dta[mdt, on = .(id==id0)], on = .(id = id1)]

    id treatment      score i.id i.treatment    i.score
1:  65         1 0.59612077   88           0 0.07473471
2:  72         1 0.14582400   94           0 0.07160314
3: 104         1 0.09878789   87           0 0.06852409

Answer 2

如果您订购数据表、制作子集并使用合并功能，我可能会有一个解决方案。 不确定它是不是最好的解决方案，但它似乎适用于我理解你想要做的事情，而且它肯定会比你的循环更快：

library(data.table) 
dta <- data.table(id = c(5,10,22,27,45,50,58,60,61,65,68,72,73,77,87,88,94,97,104,108), 
                  treatment = c(0, 0  ,0  ,0,  0, 0, 0 ,0 , 0 ,  1, 0 ,1 ,0, 0 ,0 ,0 ,0 ,0 ,1 ,0),
                  score = c(0.02381024, 0.05428605, 0.02118124, 0.01495214, 0.01877916, 0.02120360,
                            0.02207263, 0.02807019, 0.05432927, 0.59612077, 0.02482168, 0.14582400,
                            0.02371670, 0.02608826, 0.06852409, 0.07473471, 0.07160314, 0.02040747,
                            0.09878789, 0.02421807))

setkey(dta, score) # order by score
treated_nbr <- dta[treatment == 1, .N] # just to simplify the next line

selecteddata <- 
  dta[treatment == 0, 
      .SD[(.N - treated_nbr + 1):.N,
          .(correspid = id, 
            correspscore = score, 
            id = dta[treatment == 1, id])]]

这里我们取相同数量的已订购未治疗人员 (.N-treated_nbr+1):.N)，以便他们的分数与已订购人员最接近，并将 id 合并到已治疗人员的 id (id = dta[,.SD[treatment == 1,id]])

setkey(selecteddata, id)
setkey(dta, id)
selecteddata[dta] # do the merging

不确定它是否正是您想要的，因为我意识到它只有在您处理的分数高于未处理的分数时才有效（在您的示例中就是这种情况）。 您可以添加一个条件以仅针对得分高于未治疗者的治疗者使用建议的解决方案，否则执行其余操作（否则我看不到直接简单的解决方案）