EnlightenCode

二叉树遍历的统一解法 - 双色标记法 ​

背景介绍 ​

在前端开发中，树形结构随处可见（比如 DOM 树、组件树、状态树等）。理解并掌握树的遍历方法对于前端开发者来说尤为重要。本文将介绍一种优雅的统一遍历方案 - 双色标记法。

双色标记法原理 ​

这个方法的灵感来自于垃圾回收算法中的三色标记法，我们简化为双色标记：

  • 白色(WHITE)：表示节点尚未被访问
  • 灰色(GRAY)：表示节点已被访问

核心思路 ​

  1. 使用栈来维护遍历顺序
  2. 每个入栈的节点都携带一个颜色标记
  3. 通过调整入栈顺序来实现不同的遍历方式

代码实现 ​

通用模板 ​

javascript
const traversal = (root) => {
    const [WHITE, GRAY] = [0, 1];
    const result = [];
    const stack = [[WHITE, root]];
    
    while (stack.length) {
        const [color, node] = stack.pop();
        
        if (!node) continue;
        
        if (color === WHITE) {
            // 根据需要的遍历顺序调整以下三个入栈顺序
            stack.push([WHITE, node.right]);  // 右子节点
            stack.push([GRAY, node]);         // 当前节点
            stack.push([WHITE, node.left]);   // 左子节点
        } else {
            result.push(node.val);
        }
    }
    
    return result;
}

三种遍历方式的实现 ​

只需要调整入栈顺序，就能实现不同的遍历方式：

javascript
// 前序遍历 (根->左->右)
if (color === WHITE) {
    stack.push([WHITE, node.right]);
    stack.push([WHITE, node.left]);
    stack.push([GRAY, node]);     // 根节点最后入栈，最先访问
}

// 中序遍历 (左->根->右)
if (color === WHITE) {
    stack.push([WHITE, node.right]);
    stack.push([GRAY, node]);     // 根节点在中间入栈
    stack.push([WHITE, node.left]);
}

// 后序遍历 (左->右->根)
if (color === WHITE) {
    stack.push([GRAY, node]);     // 根节点最先入栈，最后访问
    stack.push([WHITE, node.right]);
    stack.push([WHITE, node.left]);
}

实际应用示例 ​

javascript
// 二叉树节点定义
class TreeNode {
    constructor(val) {
        this.val = val;
        this.left = null;
        this.right = null;
    }
}

// 创建示例树
const root = new TreeNode(1);
root.left = new TreeNode(2);
root.right = new TreeNode(3);
root.left.left = new TreeNode(4);
root.left.right = new TreeNode(5);

// 测试不同遍历方式
console.log('前序遍历:', preorderTraversal(root));  // [1, 2, 4, 5, 3]
console.log('中序遍历:', inorderTraversal(root));   // [4, 2, 5, 1, 3]
console.log('后序遍历:', postorderTraversal(root)); // [4, 5, 2, 3, 1]

优化建议 ​

对于前序遍历，其实可以使用更简单的方式实现：

javascript
const simplePreorder = (root) => {
    if (!root) return [];
    
    const result = [];
    const stack = [root];
    
    while (stack.length) {
        const node = stack.pop();
        if (!node) continue;
        
        result.push(node.val);
        // 注意：先压入右子节点，再压入左子节点
        stack.push(node.right);
        stack.push(node.left);
    }
    
    return result;
};

总结 ​

  1. 双色标记法提供了一种统一的思路来处理二叉树的三种遍历方式
  2. 通过调整入栈顺序，可以轻松实现不同的遍历顺序
  3. 这种方法在处理复杂的树形结构时特别有用，比如在开发树形组件或处理虚拟DOM时

练习建议 ​

  1. 尝试在 LeetCode 上实现这些遍历方法
  2. 将这种思想应用到前端实际项目中的树形结构处理
  3. 尝试扩展这种方法来处理 N 叉树的遍历

相关阅读 ​