使用提供的变量向量化函数

Question

我一直在尝试将一个函数从if(){} 向量化为ifelse()。当函数的所有参数都包含在它正在处理的数据集中时，它可以正常工作，但是如果我将参数作为字符串提供，则矢量化停止并且第一个结果用于整个数据集。

这是一个例子

# data
dat <- data.frame(var1 = rep(c(0,1), 4), 
                  var2 = c(rep("a", 4), rep("b", 4))
                  )

# function
my_fun <- function(x, y){
  z <- ifelse(y == "a", fun_a(x), fun_b(x))
  return(z)
}

fun_a <- function(x){
  z <- ifelse(x == 0, "zero", x)
  return(z)
}

fun_b <- function(x){
  z <- ifelse(x == 1, "ONE", x)
  return(z)
}
dat$var3 <- my_fun(dat$var1, dat$var2)

这会返回我所期望的，具有基于 var1 和 var2 的逐行值的向量

> dat
  var1 var2 var3
1    0    a zero
2    1    a    1
3    0    a zero
4    1    a    1
5    0    b    0
6    1    b  ONE
7    0    b    0
8    1    b  ONE

但是，我想在不包含 var2 的不同数据集上使用此函数。我意识到一个简单的方法是将 var2 添加为数据集中的额外列，但我真的不想这样做。

当我将 var2 作为字符串提供时会发生这种情况：

other_dat <- data.frame(var1 = rep(c(0,1), 4))
other_dat$var3 <- my_fun(other_dat$var1, y = "a")
other_dat
  var1 var3
1    0 zero
2    1 zero
3    0 zero
4    1 zero
5    0 zero
6    1 zero
7    0 zero
8    1 zero

如何向量化这个函数，让它接受一个字符串参数并返回我想要的结果？

Answer 1

您可以对 y 进行矢量化，即使 y 的长度与 x 相似，然后 ifelse 将对所有值应用函数 my_func。修改后的代码：

# data
dat <- data.frame(var1 = rep(c(0,1), 4), 
                  var2 = c(rep("a", 4), rep("b", 4))
                  )

# function
my_fun <- function(x, y){
  if(length(y) == 1) {
    y <- rep(y, length(x))
  }
  z <- ifelse(y == "a", fun_a(x), fun_b(x))
  return(z)
}

fun_a <- function(x){
  z <- ifelse(x == 0, "zero", x)
  return(z)
}

fun_b <- function(x){
  z <- ifelse(x == 1, "ONE", x)
  return(z)
}
dat$var3 <- my_fun(dat$var1, "a")

other_dat <- data.frame(var1 = rep(c(0,1), 4))
other_dat$var3 <- my_fun(other_dat$var1, y = "a")
other_dat

希望这会有所帮助。