如何将求和树递归解法转为迭代解法并维持O(H)空间复杂度

问题背景 #

要将二叉树转换为Sum Tree，规则是每个节点的新值等于其左右子树的Sum Tree之和，原始节点值需要用于计算父节点的总和。递归解法已实现，时间复杂度O(N)，空间复杂度O(H)（H为树的高度），但如果直接用常规迭代后序遍历，用HashMap存储子树和会导致空间复杂度升至O(N)，需要找到模拟递归栈的方式，保持O(H)空间。

递归解法回顾 #

递归版本利用后序遍历特性，先处理左右子树，再更新当前节点值，同时返回当前节点原始值加上子树总和（供父节点计算使用）：

class Solution{
    public void toSumTree(Node root){
         solve(root);
    }
    private int solve(Node node){
        if(node == null){
            return 0;
        }
        
        int leftSubtreeValue = solve(node.left);
        int rightSubtreeValue = solve(node.right);
        int currentNodeValue = node.data;
        
        node.data = leftSubtreeValue + rightSubtreeValue;
        
        return currentNodeValue + leftSubtreeValue + rightSubtreeValue;
    }
}

迭代版核心思路 #

递归栈的每个栈帧会保存当前节点、左右子树计算结果、节点原始值。要模拟这个过程，我们可以在栈中存储节点+访问状态：

• 状态0：节点第一次入栈，未处理左右子节点，先遍历左子树
• 状态1：节点左子树处理完成，开始处理右子树
• 状态2：节点左右子树都处理完成，可计算并更新当前节点值，同时传递总贡献给父节点

迭代实现代码 #

class Solution {
    // 自定义内部类，存储节点、访问状态及计算所需的中间值
    private static class StackFrame {
        Node node;
        int state;
        int leftSum;
        int rightSum;
        int originalValue;

        StackFrame(Node node, int state) {
            this.node = node;
            this.state = state;
        }
    }

    public void toSumTree(Node root) {
        if (root == null) {
            return;
        }

        Stack<StackFrame> stack = new Stack<>();
        stack.push(new StackFrame(root, 0));

        while (!stack.isEmpty()) {
            StackFrame frame = stack.peek();
            Node currentNode = frame.node;

            if (frame.state == 0) {
                // 第一次处理：记录原始值，标记状态后处理左子树
                frame.originalValue = currentNode.data;
                frame.state = 1;
                if (currentNode.left != null) {
                    stack.push(new StackFrame(currentNode.left, 0));
                }
            } else if (frame.state == 1) {
                // 左子树处理完成：标记状态后处理右子树
                frame.state = 2;
                if (currentNode.right != null) {
                    stack.push(new StackFrame(currentNode.right, 0));
                }
            } else {
                // 左右子树处理完成：更新节点值，传递总贡献给父节点
                currentNode.data = frame.leftSum + frame.rightSum;
                int totalContribution = frame.originalValue + frame.leftSum + frame.rightSum;
                stack.pop();
                if (!stack.isEmpty()) {
                    StackFrame parentFrame = stack.peek();
                    if (parentFrame.node.left == currentNode) {
                        parentFrame.leftSum = totalContribution;
                    } else {
                        parentFrame.rightSum = totalContribution;
                    }
                }
            }
        }
    }
}

简化版实现（双栈方式） #

如果不想自定义类，也可以用两个栈分别存储节点状态和中间计算值：

class Solution {
    public void toSumTree(Node root) {
        if (root == null) return;

        Stack<Pair<Node, Boolean>> stack = new Stack<>();
        Stack<Integer> valueStack = new Stack<>(); // 存储子树贡献值

        stack.push(new Pair<>(root, false));

        while (!stack.isEmpty()) {
            Pair<Node, Boolean> pair = stack.pop();
            Node node = pair.getKey();
            boolean isProcessed = pair.getValue();

            if (isProcessed) {
                // 处理节点：取出左右子树贡献，更新当前节点值并压入总贡献
                int rightSum = valueStack.pop();
                int leftSum = valueStack.pop();
                int originalVal = node.data;
                node.data = leftSum + rightSum;
                valueStack.push(originalVal + leftSum + rightSum);
            } else {
                // 先压回当前节点（标记为待处理），再依次压入右、左子节点
                stack.push(new Pair<>(node, true));
                if (node.right != null) {
                    stack.push(new Pair<>(node.right, false));
                }
                if (node.left != null) {
                    stack.push(new Pair<>(node.left, false));
                }
                // 压入初始0值，对应叶子节点的左右子树贡献
                valueStack.push(0);
                valueStack.push(0);
            }
        }
    }
}

复杂度分析 #

• 时间复杂度：O(N)，每个节点入栈、出栈各一次，所有操作均为常数时间
• 空间复杂度：O(H)，栈的深度等于树的高度，最坏情况（链式树）为O(N)，平衡树场景下为O(logN)

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