大型数据集下R语言高效计算种群世代数的方法问询

问题分析 #

你的代码在3万条数据上耗时数小时完全不正常，核心原因是which(Finaldf$anonID == Finaldf$ParentAnonID[current_id])这一行：每次查找父代ID时都会遍历整个数据集，时间复杂度为O(N)。如果平均每个个体有K代，总计算量就是O(N*K)，3万条数据加上多层世代，计算量会呈指数级增长，直接导致卡顿。

优化方案 #

以下几种方法可以将计算效率提升几个数量级：

方案1：预建立ID到行号的映射（基础优化） #

先一次性建立anonID到行索引的映射，把每次查找从O(N)降到O(1)：

# 预构建anonID到行号的映射向量
id_to_row <- setNames(1:nrow(Finaldf), Finaldf$anonID)

# 初始化世代数列
Finaldf$Gen <- 1L

# 遍历计算每个个体的世代数
for (i in seq_len(nrow(Finaldf))) {
  current_parent <- Finaldf$ParentAnonID[i]
  current_gen <- 1L
  
  while (current_parent != -1 && !is.na(current_parent)) {
    # 通过映射直接获取父代行号，无需全表扫描
    parent_row <- id_to_row[as.character(current_parent)]
    if (is.na(parent_row)) break
    
    current_gen <- current_gen + 1L
    current_parent <- Finaldf$ParentAnonID[parent_row]
  }
  
  Finaldf$Gen[i] <- current_gen
}

方案2：记忆化递归（避免重复计算） #

利用记忆化缓存已经计算过的世代数，避免多个后代重复计算同一个父代的世代：

# 建立ID到行号的映射
id_to_row <- setNames(1:nrow(Finaldf), Finaldf$anonID)
# 初始化缓存向量，存储已计算的世代数
gen_cache <- rep(NA_integer_, nrow(Finaldf))

get_generation <- function(id) {
  row_idx <- id_to_row[as.character(id)]
  if (is.na(row_idx)) return(1L)
  
  parent_id <- Finaldf$ParentAnonID[row_idx]
  # 根节点（父ID=-1）的世代数为1
  if (parent_id == -1 || is.na(parent_id)) {
    gen_cache[row_idx] <- 1L
    return(1L)
  }
  
  parent_row <- id_to_row[as.character(parent_id)]
  # 如果父代的世代数已缓存，直接复用
  if (!is.na(gen_cache[parent_row])) {
    gen_cache[row_idx] <- gen_cache[parent_row] + 1L
    return(gen_cache[row_idx])
  }
  
  # 递归计算父代的世代数
  gen_cache[row_idx] <- get_generation(parent_id) + 1L
  return(gen_cache[row_idx])
}

# 批量计算所有个体的世代数
Finaldf$Gen <- sapply(Finaldf$anonID, get_generation)

方案3：用data.table实现向量化操作（适合超大型数据集） #

data.table的快速匹配和分组操作能进一步提升效率：

library(data.table)
setDT(Finaldf)

# 添加行索引列
Finaldf[, row_idx := .I]
# 初始化世代数为1
Finaldf[, Gen := 1L]

# 循环更新，直到所有个体的世代数都计算完成
while (nrow(Finaldf[ParentAnonID != -1 & !is.na(ParentAnonID) & Gen == 1L, ]) > 0) {
  Finaldf[ParentAnonID != -1 & !is.na(ParentAnonID), 
          Gen := Finaldf[anonID == ParentAnonID, Gen] + 1L, 
          by = row_idx]
}

方案4：用igraph构建家谱图（适合复杂谱系） #

如果你的谱系存在分支或多个根节点，用图论工具计算路径长度更直观：

library(igraph)

# 构建父代-子代的边列表（排除根节点）
edges <- Finaldf[Finaldf$ParentAnonID != -1 & !is.na(Finaldf$ParentAnonID), 
                 c("ParentAnonID", "anonID")]
# 创建有向图
family_tree <- graph_from_data_frame(edges, directed = TRUE)

# 找到所有根节点（父ID=-1的个体）
root_nodes <- Finaldf$anonID[Finaldf$ParentAnonID == -1]

# 计算每个节点到最近根节点的路径长度，加1即为世代数
Finaldf$Gen <- sapply(Finaldf$anonID, function(node) {
  if (node %in% root_nodes) {
    1L
  } else {
    # 取到最近根节点的路径长度
    path_lengths <- shortest.paths(family_tree, from = root_nodes, to = node)
    min(path_lengths) + 1L
  }
})

效果对比 #

原代码的时间复杂度为O(N*K)，优化后的方法时间复杂度可降至O(N)或O(N log N)，3万条数据的计算时间会从小时级缩短到秒级或分钟级。

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