如何在data.table的fifelse中使用向量参数生成目标输出？

数据与需求 #

现有如下数据框，各列互斥（每行仅有一个非NA值1），需要生成output列，取值为每行非NA值所在的列名：

DF <- data.frame(
  one = c(1, NA, NA, 1, NA, NA),
  two = c(NA, 1, NA, NA, NA, 1),
  three = c(NA, NA, 1, NA, 1, NA)
)

期望输出：

output = c("one", "two", "three", "one", "three", "two")

原尝试的问题（fifelse报错） #

尝试用data.table的fifelse嵌套实现时出现报错：

with(DF, fifelse(one==1, "one", fifelse(two==1,"two", "three", na="three"), 
   na=fifelse(two==1,"two", "three", na="three")))

报错信息：

Error in fifelse(one == 1, "one", fifelse(two == 1, "two", "three", na = "three"),  : 
  Length of 'na' is 6 but must be 1

原因是*fifelse的na参数要求长度为1，不支持向量输入*，而dplyr的if_else则支持该写法：

with(DF, if_else(one==1, "one", if_else(two==1,"two", "three", missing="three"), 
   missing=if_else(two==1,"two", "three", missing="three")))

data.table的正确实现 #

推荐使用data.table的fcase函数（专为多条件判断设计，比嵌套fifelse更简洁）：

library(data.table)
# 转换为data.table格式
setDT(DF)
# 生成output列
DF[, output := fcase(
  one == 1, "one",
  two == 1, "two",
  three == 1, "three"
)]

由于各列互斥，每行只会匹配到一个条件，无需额外处理NA情况。如果需要兼容极端情况，可添加default参数指定默认值。

其他简便替代方案 #

1. dplyr简洁实现 #

使用dplyr::case_when替代嵌套if_else，代码可读性更强：

library(dplyr)
DF <- DF %>%
  mutate(output = case_when(
    one == 1 ~ "one",
    two == 1 ~ "two",
    three == 1 ~ "three"
  ))

2. R Base高效实现 #

利用max.col直接定位每行非NA值的位置，再匹配列名，无需循环或嵌套判断：

# 方法1：通过逻辑矩阵定位（通用型，适用于非NA值为任意非空值的情况）
DF$output <- colnames(DF)[max.col(!is.na(DF))]

# 方法2：利用数值特性（仅当非NA值为1时适用）
DF$output <- colnames(DF)[max.col(DF, ties.method = "first")]

两种方法都比原apply写法更高效，尤其适合大数据集。

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