【问题标题】：How do I vectorize this is_prime function in R?如何在 R 中对这个 is_prime 函数进行矢量化？
【发布时间】：2020-07-03 04:28:29
【问题描述】：

我创建了这个函数，它接受数字并根据数字是否为素数返回 TRUE 或 FALSE。

is.prime <- function(num) {
   if (num == 2) {
      TRUE
   } else if (any(num %% 2:(num-1) == 0)) {
      FALSE
   } else { 
      TRUE
   }
}

然而，这个函数只接受一个值，例如这很好用：

> is_prime(17)
[1] TRUE

如果我插入一个向量，我想查看每个元素的 TRUE 或 FALSE。例如，

> is_prime(c(17,5,10,22,109,55))
[1] TRUE
Warning messages:
1: In if (x == 1) { :
  the condition has length > 1 and only the first element will be used
2: In 2:(floor(x/2)) :
  numerical expression has 6 elements: only the first used
3: In x%%2:(floor(x/2)) :
  longer object length is not a multiple of shorter object length

这是针对第一个元素进行评估的，但我想看看

TRUE TRUE FALSE FALSE TRUE FALSE

对于向量

is_prime(c(17,5,10,22,109,55))

如何修改函数以使用相同的算法进行矢量化？

【问题讨论】：

标签： r function vectorization primes

【解决方案1】：

半向量化

可以通过以矢量化方式处理偶数（和一些其他数字）来矢量化某些函数。其余部分使用vapply 处理。

helper <- function(x) {
  for (k in seq(3, round(sqrt(x)) + 1, 2)) {
    if (x %% k == 0)
      return(FALSE)
  }
  return(TRUE)
}
is.prime <- function(v) {
  out <- rep(TRUE, length(v))
  out[v %% 2 == 0 | v %in% c(1)] <- FALSE
  out[v %in% c(2, 3, 5)] <- TRUE
  indices <- which(v > 5 && v == FALSE)
  out[indices] <- vapply(v[indices], helper, logical(1))
  return(out)
}
is.prime(c(17,5,10,22,109,55))
# [1]  TRUE  TRUE FALSE FALSE  TRUE FALSE

全矢量化

如果性能受到威胁，您可以考虑使用 `Rcpp`：

c++ 文件

#include <Rcpp.h>
#include <math.h>
using namespace Rcpp;


bool is_prime(int n) {
  if ((n == 2) || (n == 3) || (n == 5)) {
    return true;
  }
  if ((n % 2 == 0) || (n == 1)) {
    return false;
  }
  int i = 3;
  while (i < round(sqrt(n)) + 1) {
    if (n % i == 0) {
      return false;
    }
    i += 2;
  }
  return true;
}


// [[Rcpp::export]]
LogicalVector is_prime(IntegerVector v) {
  int n = v.length();
  LogicalVector out = LogicalVector(n);
  for (int i = 0; i < n; i++) {
    out[i] = is_prime(v[i]);
  }
  return out;
}

R 文件

library(Rcpp)
sourceCpp('prime_fun.cpp')  # if cpp file in same dir
is_prime(c(17,5,10,22,109,55))
# [1]  TRUE  TRUE FALSE FALSE  TRUE FALSE

【讨论】：

【解决方案2】：

从技术上讲，您不能对函数进行矢量化，但可以使用方法将is.prime 函数应用于向量。

例如，使用sapply：

sapply(c(17,5,10,22,109,55), is.prime)
#[1]  TRUE  TRUE FALSE FALSE  TRUE FALSE

或者使用Vectorize，它是mapply 的包装器：

vec.is.prime <- Vectorize(is.prime)
vec.is.prime(c(17,5,10,22,109,55))
#[1]  TRUE  TRUE FALSE FALSE  TRUE FALSE

【讨论】：

【解决方案3】：

R 是一种固有的矢量化语言，您只需要在编写函数时不要破坏它...

is.prime <- function(num){
  res<-rep(TRUE, length(num))
  for(i in 2:(max(num)-1) ){
    res[(i < num) & ((num %% i) == 0) ] <- FALSE
  }
  return(res)
}

x<-c(2, 5, 17, 109, 10, 22, 55, 93)

print("Primes:")
print(x[ is.prime(x) ] )

输出：

> source('~/.active-rstudio-document')
[1] "Primes:"
[1]   2   5  17 109

【讨论】：