大数据量下高效统计分组内数值差值≤5的实例数

我有一个包含20万+行、两列（Group和Value）的dataframe：

• Group列是字符串类型
• Value列是正整数，且每组内的Value均唯一

需求是：先按Group分组，再统计分组内每行Value与其他行Value差值≤5的实例总数。举个例子：某行Group为"A"、Value为25时，需在Group"A"的所有行中查找Value处于20-30之间的行，若存在则给Group"A"的计数加1，依次处理所有行。

我用purrr实现了正确的计数逻辑，但在20万行的数据集上运行极慢。以下是简化示例和代码：

data <- data.frame(Group = c("A", "A", "A", "B", "B", "B", "B"),
                   Value = c(1, 2, 3, 9, 15, 17, 19))

count_pairs <- function(positions) {
  combn(positions, 2, function(x) abs(x[1] - x[2]) < 5) %>% sum
}

result <- data %>%
  group_by(Group) %>%
  summarise(Count = count_pairs(Value))

运行输出：

> print(result)
# A tibble: 2 × 2
  Group Count
  <chr> <int>
1 A         3
2 B         3

请问有没有计算成本更低的实现方式？

原代码慢的核心原因是combn会生成组内所有两两元素组合，时间复杂度为O(n²)——当组内数据量较大时（比如一个组有1000行，组合数就有近50万），计算量会呈爆炸式增长。下面两种方案采用排序+区间查找的思路，时间复杂度为O(n log n)，能大幅提升处理速度。

方案一：dplyr 实现 #

利用排序后的数据特性，用findInterval快速定位符合差值条件的元素范围：

library(dplyr)

result_fast <- data %>%
  group_by(Group) %>%
  arrange(Value, .by_group = TRUE) %>%
  mutate(
    # 找到当前Value+5对应的最大索引位置
    upper_pos = findInterval(Value + 5, Value),
    # 计算当前元素能匹配到的符合条件的其他元素数量（排除自身）
    cnt = upper_pos - row_number()
  ) %>%
  summarise(Count = sum(cnt))

print(result_fast)

方案二：data.table 实现 #

对于超大数据集，data.table的性能会比dplyr更优：

library(data.table)

setDT(data)

result_dt <- data[, .(Value = sort(Value)), by = Group][,
  cnt := findInterval(Value + 5, Value) - .I,
  by = Group
][, .(Count = sum(cnt)), by = Group]

print(result_dt)

两种方案的输出结果和原代码完全一致，但处理20万行数据的速度会提升几个数量级。

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