某研究院的二叉树深度优先遍历变种的算法面试题以及答案

去了某研究院面试，被面了一道算法题，觉得有点意思，所以写下来供后人参考。

题目是这样子的：

给定二叉树，二叉树的每个节点都是一个整数值，求从叶子节点到根节点的和为某数的所有路径

例如下图中，要求叶子节点到根节点的值和为14的路径为：

3，6，5
3，7，4

这道题考的是二叉树深度优先遍历的增强版，其实现代码如下：

package cn.outofmemory;

import java.util.Stack;

public class TreeSum {
    public static void main(String[] args) {
        int expectSum = 22;

        // declare tree
        TreeNode root = new TreeNode(12, new TreeNode(6, new TreeNode(4),
                new TreeNode(7)), new TreeNode(3, new TreeNode(2),
                new TreeNode(7)));

        Stack<StackData> stack = new Stack<StackData>();
        stack.add(new StackData(root, NodeType.root));

        while (!stack.isEmpty()) {
            StackData top = stack.peek();
            TreeNode temp = top.value;

            //如果栈顶元素是叶子节点，计算看是否满足条件
            if (temp.isLeaf()) {
                int sumOfStack = getSumOfStack(stack);
                if(sumOfStack == expectSum) {
                    printStack(stack);
                }

                //弹出叶子节点
                stack.pop();

                if(top.nodeType == NodeType.left) {
                    temp = stack.peek().value;
                    //如果有右子树，将右子树放入栈顶；如果没有，则循环弹出，直到有右子树的情况
                    if(temp.getRightNode() == null) {
                        while(temp.getRightNode() == null){
                            StackData popAgain = stack.peek();
                            if (popAgain.value.getRightNode() == null) {
                                popAgain = stack.pop();
                                temp = popAgain.value;
                            }
                        }
                    } else {
                        stack.push(new StackData(temp.getRightNode(),NodeType.right));
                    }
                    continue;
                } else if (top.nodeType == NodeType.right) {
                    //如果叶子节点是右子节点，那么还要再次弹出
                    StackData popAgain = stack.pop();
                    while (popAgain.nodeType == NodeType.right) {
                        popAgain = stack.pop();
                    }
                    if(!stack.isEmpty()) {
                        StackData nowTop = stack.peek();
                        if(nowTop.value.getRightNode() != null) {
                            stack.push(new StackData(nowTop.value.getRightNode(),NodeType.right));
                        }
                    }
                } else {
                    //如果是根节点说明遍历完了，要将根节点弹出并结束循环
                    stack.pop();
                    continue;
                }

            } else {            
                //将所有的左子节点压栈
                while (temp.getLeftNode() != null) {
                    temp = temp.getLeftNode();
                    stack.add(new StackData(temp, NodeType.left));
                }           
            }
        }

    }

    private static void printStack(Stack<StackData> stack) {
        for (StackData sd : stack) {
            System.out.print("" + sd.value.getValue() + " ");
        }
        System.out.println();
    }

    private static int getSumOfStack(Stack<StackData> stack) {
        int result = 0;
        for (StackData sd : stack) {
            result += sd.value.getValue();
        }
        return result;
    }


其中TreeNode类的定义如下：

package cn.outofmemory;

public class TreeNode {

    public TreeNode() {

    }

    public TreeNode(int value) {
        this.value = value;
    }

    public TreeNode(int value, TreeNode left, TreeNode right) {
        this.value = value;
        this.setLeftNode(left);
        this.setRightNode(right);
    }

    private TreeNode left;
    private TreeNode right; 
    private int value;

    protected void setLeftNode(TreeNode left) {
        this.left = left;
    }

    public TreeNode getLeftNode() {
        return this.left;
    }

    protected void setRightNode(TreeNode right) {
        this.right = right;
    }

    public TreeNode getRightNode() {
        return this.right;
    }

    public int getValue() {
        return this.value;
    }

    public boolean isLeaf() {
        return this.left == null && this.right == null;
    }


NodeType枚举的定义如下：

package cn.outofmemory;

public enum NodeType {
    root,left,right;
}

StackData类定义如下，该类对TreeNode节点进行了封装，对于每个压栈的数据都添加了NodeType的属性，通过这个NodeType可以知道当前节点的类型。

package cn.outofmemory;

public class StackData {
    public TreeNode value;
    public NodeType nodeType;

    public StackData(TreeNode value, NodeType type) {
        this.value = value;
        this.nodeType = type;
    }
}

这道题是二叉树深度优先遍历的变种，需要灵活的利用栈。

http://outofmemory.cn/code-snippet/4247/binary-tree-deep-first-enum-test