R语言如何用循环实现按规则返回数值向量对应索引的函数

核心问题分析 #

原函数仅在初始阶段判断所有元素是否相同，剔除特殊值后未执行二次校验，且使用var(x) == 0存在边界隐患：向量长度为1、包含NA值时会返回NA，导致判断逻辑失效。

当前需求仅需一次特殊值剔除操作即可完成全部校验，不需要额外使用repeat/while循环，逻辑更简洁、执行效率更高。

优化后代码实现 #

get_indexes <- function(x, special_val) {
  # 规则(a)：原始向量所有元素完全相同，直接返回第一个索引
  if (all(x == x[1])) {
    return(1)
  }
  # 提取非特殊值对应的原始索引和元素子集
  non_special_idx <- which(x != special_val)
  filtered_x <- x[non_special_idx]
  # 校验剔除特殊值后的子集是否全部元素相同
  if (all(filtered_x == filtered_x[1])) {
    return(non_special_idx[1])
  }
  # 其余场景直接返回非特殊值的全部索引
  return(non_special_idx)
}

测试验证 #

测试用例定义：

my_vec_1 <- c(1, 2, 1, 2, 3)
my_vec_2 <- c(4, 4, 4)
my_vec_3 <- c(1, 2, 1, 4, 1)
my_vec_4 <- c(3, 3, 3, 4) 
special_treatment_value <- 4

测试输出：

get_indexes(my_vec_1, special_treatment_value)
#> [1] 1 2 3 4 5 符合规则(c)预期
get_indexes(my_vec_2, special_treatment_value)
#> [1] 1 符合规则(a)预期
get_indexes(my_vec_3, special_treatment_value)
#> [1] 1 2 3 5 符合规则(b)预期
get_indexes(my_vec_4, special_treatment_value)
#> [1] 1 符合优化后预期

如果后续需要适配多轮不同特殊值剔除的扩展场景，可以将上述逻辑封装为递归逻辑，或增加循环控制参数即可。

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