首页 > 技术文章 > R 函数向量化

impw 2020-06-09 21:51 原文

自定义的函数,如果其中的计算都是向量化的, 那么函数自动地可以接受向量作为输入,结果输出向量。 比如,将每个元素都变成原来的平方的函数:

f <- function(x){
  x^2
}

如果输入一个向量,结果也是向量,输出的每个元素是输入的对应元素的相应的平方值。

但是,如下的分段函数:
g(x) = \begin{cases}
x^2 , & |x| \leq 1, \
1, & |x| > 1
\end{cases}
其一元函数版本可以写成

g <- function(x){
  if(abs(x) <= 1) {
    y <- x^2
  } else {
    y <- 1
  }
  
  y
}

但是这个函数不能处理向量输入,因为if语句的条件必须是标量条件。 一个容易想到的修改是

gv <- function(x){
  y <- numeric(length(x))
  sele <- abs(x) <= 1
  y[sele] <- x[sele]^2
  y[!sele] <- 1.0
  
  y
}

或者

gv <- function(x){
  ifelse(abs(x) <= 1, x^2, 1)
}

对于没有这样简单做法的问题,可以将原来的逻辑包在循环中,如

gv <- function(x){
  y <- numeric(length(x))
  for(i in seq(along=x)){
    if(abs(x[i]) <= 1) {
      y[i] <- x[i]^2
    } else {
      y[i] <- 1
    }
  }
  
  y
}

函数Vectorize可以将这样的操作自动化。如

g <- function(x){
  if(abs(x) <= 1) {
    y <- x^2
  } else {
    y <- 1
  }
  
  y
}
gv <- Vectorize(g)
gv(c(-2, -0.5, 0, 0.5, 1, 1.5))
## [1] 1.00 0.25 0.00 0.25 1.00 1.00

推荐阅读