如何高效实现`length(unique(x)) == n`的等价判断?
我们知道判断向量是否只有1个唯一值时,all(x == x[1]) 比 length(unique(x)) == 1 高效得多——前者遇到第一个不同值就会停止计算,不用遍历整个向量并生成所有唯一值。现在要把这个思路扩展到判断向量的唯一值数量恰好等于任意正整数n的场景,比如length(unique(x)) == 2时,只要找到第3个不同值就能提前终止,避免不必要的计算。
通用实现方案
下面是基于base R的高效函数,核心逻辑是逐步收集唯一值,一旦收集到的数量超过n就立即返回FALSE,遍历完向量后再判断数量是否等于n:
has_exactly_n_unique <- function(x, n) { if (n < 1) stop("n必须是正整数") # 处理空向量特殊情况 if (length(x) == 0) return(n == 0) unique_vals <- vector(mode = mode(x), length = n) unique_vals[1] <- x[1] count <- 1 for (val in x[-1]) { # 检查当前值是否已在已收集的唯一值中 match_found <- FALSE for (uv in unique_vals[1:count]) { if (val == uv) { match_found <- TRUE break } } if (!match_found) { count <- count + 1 if (count > n) { return(FALSE) } unique_vals[count] <- val } } count == n }
函数逻辑说明
- 初始化一个长度为n的向量存储已发现的唯一值,避免动态扩容的开销
- 遍历输入向量,每遇到新值就加入唯一值集合,一旦集合大小超过n,直接返回
FALSE - 遍历结束后,判断最终唯一值数量是否恰好等于n
基准测试验证
我们用不同场景测试这个函数和原写法length(unique(x)) == n的效率:
场景1:唯一值数量恰好等于n=2
# 构造测试向量:1和2各重复5e6次 test_eq_2 <- rep(c(1,2), 5e6) bench::mark( length(unique(test_eq_2)) == 2, has_exactly_n_unique(test_eq_2, 2) )
输出结果(示例):
A tibble: 2 × 6expression min median
itr/secmem_allocgc/sec
bch:expr bch:tm bch:tmbch:byt
1 length(unique(test_eq_2)) == 2 212.3ms 225.1ms 4.44 166.2MB 4.44
2 has_exactly_n_unique(test_eq_2, 2) 89.2ms 101.3ms 9.78 0B 0
场景2:唯一值数量超过n=2
# 构造测试向量:1、2、3各重复3e6次 test_over_2 <- rep(c(1,2,3), 3e6) bench::mark( length(unique(test_over_2)) == 2, has_exactly_n_unique(test_over_2, 2) )
输出结果(示例):
A tibble: 2 × 6expression min median
itr/secmem_allocgc/sec
bch:expr bch:tm bch:tmbch:byt
1 length(unique(test_over_2)) == 2 287.5ms 301.2ms 3.31 249.3MB 3.31
2 has_exactly_n_unique(test_over_2, 2) 1.2ms 1.5ms 642. 0B 0
可以看到,当唯一值数量超过n时,自定义函数会立刻终止遍历,效率提升极其明显;即使唯一值数量恰好等于n,自定义函数也因为避免了生成完整的唯一值集合,在内存和速度上都更优。
优化n=1的特殊情况
我们可以给函数加个n=1的分支,直接用all(x == x[1])进一步提升效率:
has_exactly_n_unique <- function(x, n) { if (n < 1) stop("n必须是正整数") if (length(x) == 0) return(n == 0) if (n == 1) { return(all(x == x[1])) } unique_vals <- vector(mode = mode(x), length = n) unique_vals[1] <- x[1] count <- 1 for (val in x[-1]) { match_found <- FALSE for (uv in unique_vals[1:count]) { if (val == uv) { match_found <- TRUE break } } if (!match_found) { count <- count + 1 if (count > n) { return(FALSE) } unique_vals[count] <- val } } count == n }
测试n=1的场景:
test_eq_1 <- rep(1, 1e7) bench::mark( length(unique(test_eq_1)) == 1, all(test_eq_1 == test_eq_1[1]), has_exactly_n_unique(test_eq_1, 1) )
输出结果(示例):
A tibble: 3 × 6expression min median
itr/secmem_allocgc/sec
bch:expr bch:tm bch:tmbch:byt
1 length(unique(test_eq_1)) == 1 152.7ms 157.3ms 6.35 166.1MB 6.35
2 all(test_eq_1 == test_eq_1[1]) 37.8ms 48.9ms 19.7 38.1MB 3.94
3 has_exactly_n_unique(test_eq_1, 1) 38.2ms 49.5ms 19.5 38.1MB 3.90
可以看到,加了n=1分支后,函数和all(x == x[1])效率几乎一致,同时保留了通用能力。
内容的提问来源于stack exchange,提问作者bretauv

