查找可遍历数组并返回符合条件元素索引的现成R函数

Base R原生实现 #

通用场景最简写法 #

绝大多数场景下直接使用which()搭配向量化条件判断即可，代码简洁且底层为C实现，性能远高于自行遍历：

y <- c(12, 44, 8, -2, 5)
which(y < 10)
# 输出：[1] 3 4 5

超大数据低内存写法 #

如果待处理数组体量极大、符合条件的样本占比极低，需要避免生成全量中间布尔数组，可以使用Filter()迭代过滤索引：

y <- c(12, 44, 8, -2, 5)
Filter(function(i) y[i] < 10, seq_along(y))
# 输出：[1] 3 4 5

该方案只会在迭代中逐个判断索引，不会生成全量逻辑数组，内存占用随符合条件的样本量增长，而非随原始数组长度增长。

purrr包实现 #

purrr中可以使用keep()函数实现迭代过滤，写法更符合函数式编程风格：

library(purrr)
y <- c(12, 44, 8, -2, 5)
seq_along(y) %>% keep(~ y[.x] < 10)
# 输出：[1] 3 4 5

该方案同样不会生成全量中间数组，内存效率与Filter()一致。

笛卡尔积场景高性能方案 #

如果待处理数据为两个数组的笛卡尔积，推荐使用data.table的CJ()函数边生成笛卡尔积边筛选，无需生成全量笛卡尔积数组，性能远高于R层遍历：

library(data.table)
# 示例两个待计算笛卡尔积的数组
arr1 <- c(3, 15, 9, 22)
arr2 <- c(7, 4, 18, 11)
# 直接返回符合条件的索引对，不会加载全量笛卡尔积到内存
CJ(idx1 = seq_along(arr1), idx2 = seq_along(arr2))[arr1[idx1] + arr2[idx2] < 20]

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