前言

在开发 UI 组件，如侧栏菜单、级联选择器、表格、树形选择器等的过程中，我们经常会遇到需要处理数据结构——Tree的情况，但是每次写起来却觉得没那么简单或者代码看起来不够优雅。这里，我就结合一些源码谈谈常见的树形结构处理的一些技巧。

基础知识

DFS（deep-first-search 深度优先遍历）和 BFS（breath-first-search 广度优先遍历）都是常见的搜索算法。往往和递归、遍历、回溯的概念相关联。在树形结构的处理中，我们往往结合这两种算法处理。 我们来看一个经典的树形结构并思考以下问题：

const data = {
  value: 0,
  children: [
    {
      value: 1,
      children: [
        { value: 11 },
        { value: 12 },
      ],
    },
    {
      value: 2,
      children: [
        { value: 21 },
      ],
    },
  ],
};

// [0, 1, 11, 12, 2, 21]

// [0, 1, 2, 11, 12, 21]

深度优先遍历

深度优先即从根节点向下遍历，一条路径走完再查找另一条路径，通常可以结合“栈”这个数据结构进行处理。

• 迭代解法

function dfs_iterate(tree, callback) {
  let stack = [...tree];
  while (stack.length) {
    let node = stack.pop();
    callback(node.value);
    if (node.children) {
      // 思考，先进后出，要保证列表顺序，需要反转数组
      stack.push(...node.children.reverse());
    }
  }
}

• 递归解法

function dfs(data, callback) {
  // 递归是天然的堆栈
  data.forEach((node) => {
    // 前序 or 后序 取决于这里先遍历根还是孩子节点
    callback(node.value);
    if (node.children) {
      dfs(node.children, callback);
    }
  });
}

广度优先遍历

广度优先即从根节点向下遍历，先找到所有最近的子节点，再向下查找，通常可以结合“队列”这个数据结构进行处理。

• 迭代解法

function bfs_iterate(tree, callback) {
  let queue = [...tree];
  while (queue.length) {
    let node = queue.shift();
    callback(node.value);
    if (node.children) {
      queue.push(...node.children);
    }
  }
}

• 递归解法

通常我们不采用递归进行广度优先遍历，因为递归是天然的栈，但是这里我们可以通过一个 level 参数标记当前深度转化为深度递归遍历，相当于层次遍历

function bfs(tree) {
  let result = [];
  dfs(tree, 0);
  function dfs(tree, level) {
    tree.forEach((node) => {
      if (result.length <= level) {
        result.push([]);
      }
      result[level].push(node.value);
      if (node.children) {
        dfs(node.children, level + 1);
      }
    });
  }
  return result;
}

扁平数组转树形结构

这里我们先通过一个例子来简单分析一下：

const data = [
   { id: 'node-1', parent: 'root' },
   { id: 'node-2', parent: 'root' },
   { id: 'node-3', parent: 'node-2' },
   { id: 'node-4', parent: 'node-2' },
   { id: 'node-5', parent: 'node-4' },
 ]

想想构建这棵树我们要怎么做呢？最容易理解的就是至上而下构建的过程，很直观的：

• 我们可以推测，最终构建的树的根节点是id === node-1 || node-2
• node-1 没有子节点
• node-2 的子节点是id === node-3 || node-4
• node-3 没有子节点
• node-4 的子节点是id === node-5
• node-5 没有子节点

最终，我们可以得到这棵树的结果是：

[
  { id: "node-1", parent: "root" },
  {
    id: "node-2",
    parent: "root",
    children: [
      { id: "node-3", parent: "node-2" },
      {
        id: "node-4",
        parent: "node-2",
        children: [{ id: "node-5", parent: "node-4" }],
      },
    ],
  },
];

理论

那么，我们如何将这个构建的过程抽象化为可执行的步骤呢？这里我们将树构建的过程拆分为三个阶段：

• 找到所有父-子关系
• 确定所有根节点
• 以根节点为起点，递归构建树

代码实战

找到所有父-子关系

• 根据指定的idKey 与 parentKey
• 当满足A[parentKey] === B[idKey]时，对象 A 是对象 B 的 child

const findParentChildren = (items, idKey, parentKey) => {
  let parnetWrapper = new Map();
  items.forEach((item) => {
    if (parnetWrapper.has(item[parentKey])) {
      let children = parnetWrapper.get(item[parentKey]);
      children.push(item);
    } else {
      parnetWrapper.set(item[parentKey], [item]);
    }
  });
  return parnetWrapper;
};

在上面例子中

const idKey = "id";
const parentKey = "parent";
const parnetWrapper = findParentChildren(data);

确定所有根节点

根节点的确定可以分两种情况讨论：

• roolVal 已知，根据parentKey === rootVal可以确定哪些节点是根节点
• roolVal 未知，只要parentKey 不在任意 idKey中，该节点就是根节点

const findRoots = (items) => {
  if (rootVal) {
    return items.filter((ele) => ele[parentKey] === rootVal);
  } else {
    return items.filter(
      (parent) => !items.find((item) => item[idKey] === parent[parentKey])
    );
  }
};

在上面例子中，我们同样可以用两种方式确定：

• 已知rootVal = "root"
• 无法找到id === "root"的数据，因此id === "node-1" || "node-2"的节点是根节点

// const rootVal = "root";
const roots = findRoots(data);

递归构建树

const buildTree = (roots, wrapper) => {
  return roots.map((item) => {
    if (wrapper.has(item[idKey])) {
      return {
        ...item,
        children: buildTree(wrapper.get(item[idKey]), wrapper),
      };
    } else {
      return item;
    }
  });
};

TypeScript 实现版本见 buildTree

树形结构转扁平化数组

在组件开发中,扁平化数组也是很常见的操作。比如，级联组件一开始传入一个树形结构，为了方便后面的过滤和搜索，我们会先对树形结构进行扁平化处理。这个过程其实就是数组遍历的过程，递归或者迭代都可以进行处理。

const flattenTree = (
  items,
  childrenKey = "children"
) => {
  const flattenOptions = [];
  dfs(items);
  return flattenOptions;
  function dfs(nodes) {
    if (nodes == null) {
      return;
    }
    for (const node of nodes) {
      if (!node[childrenKey] || !node[childrenKey].length) {
        flattenOptions.push(node);
      } else {
        dfs(node[childrenKey]);
      }
    }
  }
};

TypeScript 实现版本见 flattenTree

参考源码