在R中基于节点信息自动生成二叉拆分对应边列表数据框的方法

实现逻辑 #

通过栈模拟深度优先的左子树优先挂载规则，适配你提到的二叉拆分要求：

• 每次优先给当前最深层的未饱和（子节点数<2）非叶节点挂载子节点
• 非叶节点挂载完2个子节点后自动回退到上层节点继续挂载

完整R实现代码 #

generate_edges <- function(df) {
  # 初始化结果、栈、子节点计数表
  edges <- data.frame(from = integer(), to = integer())
  stack <- c(1) # 根节点入栈
  child_count <- setNames(rep(0, max(df$node)), df$node)
  
  # 从第2个节点开始遍历挂载
  for (current_to in 2:nrow(df)) {
    # 取栈顶作为当前父节点
    current_from <- tail(stack, 1)
    # 新增边记录
    edges <- rbind(edges, data.frame(from = current_from, to = current_to))
    # 父节点子节点计数+1
    child_count[as.character(current_from)] <- child_count[as.character(current_from)] + 1
    # 如果当前挂载的节点不是叶节点，压入栈作为下一级父节点
    if (!df$isLeaf[df$node == current_to]) {
      stack <- c(stack, current_to)
    }
    # 检查栈顶节点是否已经满负载（2个子节点），满负载则弹出栈
    while (length(stack) > 0 && child_count[as.character(tail(stack, 1))] >= 2) {
      stack <- stack[-length(stack)]
    }
  }
  rownames(edges) <- NULL
  return(edges)
}

测试验证 #

测试用例1（小规模） #

dfTest1 <- data.frame(
  node = c(1,2,3,4,5),
  var = c("milk", NA, "coffee", NA, NA),
  isLeaf = c(F, T, F, T, T)
)
edges1 <- generate_edges(dfTest1)
print(edges1)

输出结果与预期一致：

from to
1    1  2
2    1  3
3    3  4
4    3  5

测试用例2（大规模） #

dfTest2 <- data.frame(
  node = c(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11),
  var = c("milk", "milk", NA, NA, "coffee", "sugar", NA, NA, "milk", NA, NA),
  isLeaf = c(F, F, T, T, F, F, T, T, F, T, T)
)
edges2 <- generate_edges(dfTest2)
print(edges2)

输出结果与预期一致：

from to
1     1  2
2     2  3
3     2  4
4     1  5
5     5  6
6     6  7
7     6  8
8     5  9
9     9 10
10    9 11

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