如何将扁平数组转换为指定最大深度的树形结构？

原始数据 #

[
  {
    id: "1",
    parent: null
  },
  {
    id: "2",
    parent: null
  },
  {
    id: "3",
    parent: "1"
  },
  {
    id: "4",
    parent: "3"
  },
]

需求说明 #

常规数组转树形结构的工具会生成完全嵌套的树形结构，而我们需要限制树形的最大深度：当节点层级超过指定深度时，不再继续嵌套，直接将其挂载到对应层级的祖先节点的children列表中。

常规全嵌套结果 #

[
  {
    data: { id: "1", parent: null },
    children: [
      {
        data: { id: "3", parent: "1" },
        children: [
          { data: { id: "4", parent: "3" } }
        ] 
      }
    ]
  },
  {
    data: { id: "2", parent: null }
  }
]

指定深度为1的期望结果 #

[
  {
    data: { id: "1", parent: null },
    children: [
      { data: { id: "3", parent: "1" } },
      { data: { id: "4", parent: "3" } }
    ]
  },
  {
    data: { id: "2", parent: null }
  }
]

实现思路与代码示例 #

无需先构建全嵌套树再扁平化，直接在构建过程中控制深度，步骤如下：

• 建立节点ID到节点的映射，同时初始化每个节点的结构
• 计算每个节点的层级深度（根节点深度为0）
• 根据节点深度判断挂载目标：深度未超限则挂载到原始父节点；深度超限则找到层级等于最大深度的祖先节点，挂载到该节点下

function arrayToTreeWithMaxDepth(items, maxDepth) {
  const nodeMap = new Map();
  const rootNodes = [];
  const depthMap = new Map();

  // 初始化节点并计算每个节点的深度
  items.forEach(item => {
    const node = { data: item, children: [] };
    nodeMap.set(item.id, node);

    if (item.parent === null) {
      depthMap.set(item.id, 0);
      rootNodes.push(node);
    } else {
      const parentDepth = depthMap.get(item.parent);
      depthMap.set(item.id, parentDepth + 1);
    }
  });

  // 将节点挂载到对应层级的父节点下
  items.forEach(item => {
    const node = nodeMap.get(item.id);
    const nodeDepth = depthMap.get(item.id);

    if (nodeDepth === 0) return;

    let targetParentId;
    if (nodeDepth <= maxDepth) {
      targetParentId = item.parent;
    } else {
      // 回溯找到深度为maxDepth的祖先节点
      let currentId = item.parent;
      while (depthMap.get(currentId) > maxDepth) {
        currentId = nodeMap.get(currentId).data.parent;
      }
      targetParentId = currentId;
    }

    nodeMap.get(targetParentId).children.push(node);
  });

  return rootNodes;
}

// 测试示例
const items = [
  { id: "1", parent: null },
  { id: "2", parent: null },
  { id: "3", parent: "1" },
  { id: "4", parent: "3" },
];

// 指定最大深度为1
const result = arrayToTreeWithMaxDepth(items, 1);
console.log(JSON.stringify(result, null, 2));

补充说明 #

• 该方法仅需遍历数组两次，效率较高，且保留节点的原始输入顺序
• 默认输入节点的父节点均存在于数组中，若需处理无效父节点，可添加额外的错误逻辑
• 可根据需求调整节点的包装结构（例如移除data字段，直接使用原始节点格式）

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