获取与值匹配的 data.table 列的索引答案

【问题标题】：Get index of data.table column that matches a value获取与值匹配的 data.table 列的索引
【发布时间】：2019-11-07 13:47:15
【问题描述】：

DT = data.table(
      id = 1:5,
      a  = c(0,1,0,2,5),
      b  = c(1,0,2,4,4),
      c  = c(1,2,0,0,5))

#     id a b c
# 1:  1  0 1 1
# 2:  2  1 0 2
# 3:  3  0 2 0
# 4:  4  2 4 0
# 5:  5  5 4 5

我想识别左边第一列有 0，并将列索引放在idx。

#     id a b c idx
# 1:  1  0 1 1 2 
# 2:  2  1 0 2 3
# 3:  3  0 2 0 2
# 4:  4  2 4 0 4
# 5:  5  5 4 5 NA

（也欢迎使用非data.table 解决方案，例如使用dplyr）

【问题讨论】：

标签： r dplyr data.table

【解决方案1】：

基于 R 的想法可以是，

replace(max.col(-DT[, -1] == 0, ties.method = 'first') + 1, rowSums(DT == 0) == 0, NA)

#or break it into two lines If you want,
i1 <- max.col(-DT[,-1] == 0, ties.method = 'first') + 1
replace(i1, rowSums(DT == 0) == 0, NA)

#[1]  2  3  2  4 NA

【讨论】：

【解决方案2】：

data.table 的解决方案可能是：

DT[, idx := apply(.SD, 1, function(x) first(which(x == 0 )))]

#   id a b c idx
#1:  1 0 1 1   2
#2:  2 1 0 2   3
#3:  3 0 2 0   2
#4:  4 2 4 0   4
#5:  5 5 4 5  NA

【讨论】：

【解决方案3】：

dplyr 的一种可能是：

DT %>%
 mutate(idx = if_else(rowSums(. == 0) == 0,
                      NA_integer_,
                      max.col(- ., ties.method = "first")))

  id a b c idx
1  1 0 1 1   2
2  2 1 0 2   3
3  3 0 2 0   2
4  4 2 4 0   4
5  5 5 4 5  NA

data.table 也一样：

DT[, idx := ifelse(rowSums(.SD == 0) == 0,
                   NA_integer_,
                   max.col(- .SD, ties.method = "first"))]

【讨论】：

哦没看到你。你为什么称它为dplyr 可能性？实际上，为什么要加载 dplyr 以使用来自基础 R 的函数来执行 data.table 格式的操作？
OP 提到他/她也可能对某些dplyr 的可能性感兴趣。它可以很容易地重写为base R 格式（就像你已经做过的那样），但是如果有人坚持用dplyr 这样做，那么他/她可以使用它:)
唯一的 dplyr 部分是管道（所以它只是 magrittr？）。如果我们想使用 dplyr (tidyverse) 那么使用相关函数，例如if_else 而不是ifelse?
@zx8754 mutate() 来自dplyr。使用if_else() 是个好点，相应地更新它。

【解决方案4】：

DT[, idx := which(.SD == 0)[1] + 1L, by = id]
DT
#    id a b c idx
# 1:  1 0 1 1   2
# 2:  2 1 0 2   3
# 3:  3 0 2 0   2
# 4:  4 2 4 0   4
# 5:  5 5 4 5  NA

【讨论】：

【解决方案5】：

评论太长，因此需要 comm wiki。类似于 Sotos 和 tmfmnk，如果在id 中找不到值，则可以避开rowSums：

DT[, idx := {
    x <- max.col(.SD==0, "first")
    replace(x, x==1L, NA_integer_)
}]

计时码：

library(data.table)
set.seed(0L)
nr <- 1e7
DT <- data.table(id=1:nr, a=sample(0:5, nr, TRUE), b=sample(0:5, nr, TRUE), c=sample(0:5, nr, TRUE))
DT0 <- copy(DT)
DT1 <- copy(DT)
DT2 <- copy(DT)
DT3 <- copy(DT)
DT4 <- copy(DT)

mtd0 <- function() {
    replace(max.col(-DT0[, -1] == 0, ties.method = 'first') + 1, rowSums(DT == 0) == 0, NA)

    #or break it into two lines If you want,
    i1 <- max.col(-DT0[,-1] == 0, ties.method = 'first') + 1
    replace(i1, rowSums(DT0 == 0) == 0, NA)
}

mtd1 <- function() {
    DT1[, idx := apply(.SD, 1, function(x) first(which(x == 0 )))]
}

mtd2 <- function() {
    DT2[, idx := which(.SD == 0)[1], by = id]
}

mtd3 <- function() {
    DT2[, idx := fifelse(rowSums(.SD == 0) == 0,
        NA_integer_,
        max.col(-.SD, ties.method = "first"))]
}

mtd4 <- function() {
    DT4[, idx := {
        x <- max.col(.SD==0, "first")
        replace(x, x==1L, NA_integer_)
    }]
}

bench::mark(mtd0(), mtd1(), mtd2(), 
    mtd3(), mtd4(), check=FALSE)

时间安排：

# A tibble: 5 x 13
  expression      min   median `itr/sec` mem_alloc `gc/sec` n_itr  n_gc total_time result    memory    time   gc    
  <bch:expr> <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl> <int> <dbl>   <bch:tm> <list>    <list>    <list> <list>
1 mtd0()     136.98ms 145.67ms    5.41     347.8MB    1.80      3     1    554.1ms <dbl [1,~ <df[,3] ~ <bch:~ <tibb~
2 mtd1()        8.66s    8.66s    0.115     76.3MB    0.346     1     3      8.66s <df[,5] ~ <df[,3] ~ <bch:~ <tibb~
3 mtd2()        57.3s    57.3s    0.0175    22.6MB    0.436     1    25      57.3s <df[,5] ~ <df[,3] ~ <bch:~ <tibb~
4 mtd3()       86.6ms  90.69ms   10.3      171.7MB    1.72      6     1   582.21ms <df[,5] ~ <df[,3] ~ <bch:~ <tibb~
5 mtd4()       48.9ms  50.12ms   18.4       97.6MB    1.84     10     1   544.43ms <df[,5] ~ <df[,3] ~ <bch:~ <tibb~

【讨论】：