使用data.table分配多个相互引用的变量

4

我希望能够给一个变量赋值,然后使用该变量创建一个新的变量。 data.table 的语法支持多重赋值,但显然不能用于内部引用。 在我的实际用例中,“i”和“by”子句更加复杂,因此我不想像这样重复编写代码:

require(data.table)

dt <- data.table(
  x = 1:5, 
  y = 2:6
)

# this works
dt[x == 3, z1 := x + y]
dt[x == 3, z2 := z1 + 5]

# but I wish this worked
dt[x == 3, `:=`(
  z1 = x + y,
  z2 = z1 + 5
)]

相比之下,dplyr 可以实现这个功能:
require(dplyr)

df <- data.frame(
  x = 1:5, 
  y = 2:6
)

df <- mutate(df,
  z1 = x + y,
  z2 = z1 + 5
)

有没有一种使用data.table的简洁方法来做到这一点?

编辑: 稍微调整akrun的解决方案,我找到了一种保持易读的连续语法的方法。只需在列表外执行所有操作:

dt[x==3, c('z1','z2','z3') := {
  z1 <- x+y
  z2 <- z1 + 5
  z3 <- z2 + 6
  list(z1, z2, z3) 
}]

1
你展示的 dplyr 选项与 data.table 不同,因为它没有过滤 x==3。在 dplyr 中,我猜我们需要使用 ifelse 或进行 filter,然后进行 mutateleft_join,如果我没错的话,这应该是比较耗费资源的。 - akrun
2
没错。我只是举了一个内部参考的例子。我想我会保持原样,因为我已经知道dplyr比较慢了。谢谢你的帮助。 - rsoren
1个回答

4
我们可以使用花括号创建临时变量,然后将它们与基于该变量的计算结果一起放在列表中,并将 (:=) 分配给我们需要创建的列。
dt[x==3, c('z1', 'z2') := {
             z1 <- x+y
             list(z1, z1+5) 
             }]
dt
#   x y z1 z2
#1: 1 2 NA NA
#2: 2 3 NA NA
#3: 3 4  7 12
#4: 4 5 NA NA
#5: 5 6 NA NA

为了让它变得更快,我们可以使用 setkey
setkey(dt, x)[(3),  c('z1', 'z2') := {
                                   z1 <- x+y
                              list(z1, z1+5)
                  }]

基准测试

set.seed(24)
dt1 <- data.table(x = sample(1:9, 1e8, replace=TRUE), y = sample(5:9, 1e8, replace=TRUE))

dt2 <- copy(dt1)
dt3 <- copy(dt1)

akrun1 <- function(){dt1[x==3, c('z1', 'z2') := {
             z1 <- x+y
                 list(z1, z1+5) 
             }]
   }

akrun2 <- function() {setkey(dt3, x)[(3),  c('z1', 'z2') := {
                                   z1 <- x+y
                              list(z1, z1+5)
                  }]
}


rsoren  <- function() {
    dt2[x == 3, z1 := x + y]
    dt2[x == 3, z2 := z1 + 5]
        }



library(microbenchmark)
microbenchmark(akrun1(), akrun2(), rsoren(), unit= "relative", times = 20L)
#Unit: relative
#     expr      min       lq     mean   median       uq       max neval
# akrun1() 1.597267 1.605404 1.393016 1.642584 1.538929 0.8634406    20
# akrun2() 1.000000 1.000000 1.000000 1.000000 1.000000 1.0000000    20
# rsoren() 2.584153 2.586185 2.179601 2.694469 2.468219 0.9740701    20

@rsoren,这里只发生一次赋值(:=),如果你已经检查过代码,也不必多次执行x==3,否则会再次减慢代码。 - akrun
@rsoren 我的意思是如果你仔细查看了代码 - akrun
谢谢。我会保持这个开放一段时间,希望有人能避免使用 c('var1','var2','var3') := {stuff1;stuff2;stuff3} 这样的语法,因为我认为它不如我的 dplyr 示例易读。 - rsoren
@rsoren 如果我引用你的话“我宁愿不要有这样重复的代码:”。那么,你是在寻求效率还是重复代码? - akrun
@rsoren,我的两种方法似乎比你的更快,并且还避免了重复的语法。已更新基准测试结果。 - akrun
显示剩余5条评论

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接