代码改变世界

在二叉树中找出和为某一值的所有路径

2017-02-11 21:25  微服务专家  阅读(344)  评论(0编辑  收藏  举报
import java.util.*;


public class Main {
    public static void main(String[] args) {
        Scanner cin = new Scanner(System.in);
        long expectedSum = 0;
        boolean gotSum = false;
        while (cin.hasNext() && !gotSum) {
            String firstLine = cin.nextLine();
            if (firstLine == null || firstLine.isEmpty()) {
                continue;
            }
            if (firstLine.trim().isEmpty()) {
                continue;
            }
            try {
                expectedSum = Long.parseLong(firstLine);
                gotSum = true;
                break;
            } catch (Exception e) {
                continue;
            }
        }
        TreeNode treeRoot = null;
        boolean gotNodes = false;
        while (cin.hasNext() && !gotNodes) {
            String nodeLine = cin.nextLine();
            if (nodeLine == null || nodeLine.isEmpty()) {
                continue;
            }
            if (nodeLine.trim().isEmpty()) {
                continue;
            }
            try {
                String nodeStrings[] = nodeLine.split(",");
                treeRoot = createTree(nodeStrings);
                gotNodes = true;
                break;
            } catch (Exception e) {
                continue;
            }
        }

        List<Stack<TraceStep>> traces = findTrace(treeRoot, expectedSum);
//
        if (traces.isEmpty()) {
            System.out.println("error");
            return;
        }
        for (Stack<TraceStep> trace : traces) {
            StringBuffer stringBuffer = new StringBuffer();
            for (int i = 0; i < trace.size() - 1; i++) {
                stringBuffer.append(trace.get(i).getNode().getValue()).append(',');
            }
            stringBuffer.append(trace.pop());
            System.out.println(stringBuffer.toString());
        }

    }

    private static TreeNode createTree(String nodeStrings[]) {
        TreeNode root = new TreeNode(Long.parseLong(nodeStrings[0]));
        if (nodeStrings.length > 1) {
            addChildRen(new TreeNode[]{root}, Arrays.copyOfRange(nodeStrings, 1, nodeStrings.length));
        }

        return root;
    }

    private static void addChildRen(TreeNode[] leafs, String nodeStrings[]) {

        if (nodeStrings.length > leafs.length * 2) {
            TreeNode[] newLeafs = new TreeNode[leafs.length * 2];
            for (int i = 0; i < newLeafs.length; i++) {
                newLeafs[i] = new TreeNode(Long.parseLong(nodeStrings[i]));

            }
            for (int i = 0; i < newLeafs.length; i++) {
                if (i % 2 == 0) {
                    leafs[i / 2].setLeft(newLeafs[i]);
                }
                if (i % 2 == 1) {
                    leafs[i / 2].setRight(newLeafs[i]);
                }
            }
            addChildRen(newLeafs, Arrays.copyOfRange(nodeStrings, newLeafs.length, nodeStrings.length));
        } else {
            TreeNode[] newLeafs = new TreeNode[nodeStrings.length];
            for (int i = 0; i < newLeafs.length; i++) {
                newLeafs[i] = new TreeNode(Long.parseLong(nodeStrings[i]));

            }
            for (int i = 0; i < newLeafs.length; i++) {
                if (i % 2 == 0) {
                    leafs[i / 2].setLeft(newLeafs[i]);
                }
                if (i % 2 == 1) {
                    leafs[i / 2].setRight(newLeafs[i]);
                }
            }
        }
    }

    public static List<Stack<TraceStep>> findTrace(TreeNode treeRoot, long sumExpected) {
        List<Stack<TraceStep>> allTraces = new ArrayList<Stack<TraceStep>>();
        Stack<TraceStep> trace = new Stack<TraceStep>();

        TraceStep current = new TraceStep(treeRoot);
        trace.add(current);
        while (!trace.isEmpty()) {
            if (!current.isLeftVisited()) {
                current.setLeftVisited(true);
                if (current.getNode().getLeft() != null) {
                    trace.add(new TraceStep(current.getNode().getLeft()));
                    current = trace.peek();

                }

            } else if (!current.isRightVisited()) {
                current.setRightVisited(true);
                if (current.getNode().getRight() != null) {
                    trace.add(new TraceStep(current.getNode().getRight()));
                    current = trace.peek();
                }
            } else {
                if ((current.getNode().getLeft() == null) && (current.getNode().getRight() == null)) {
                    if (sum(trace) == sumExpected) {
                        allTraces.add(copy(trace));
                    }
                }
                trace.pop();
                if (!trace.isEmpty()) {
                    current = trace.peek();
                }

            }

        }


        return allTraces;
    }

    private static Stack<TraceStep> copy(Stack<TraceStep> trace) {
        Stack<TraceStep> aCopy = new Stack<TraceStep>();
        for (TraceStep traceStep : trace) {
            aCopy.add(new TraceStep(traceStep));
        }
        return aCopy;
    }

    private static long sum(Stack<TraceStep> trace) {
        long result = 0;
        for (TraceStep traceStep : trace) {
            result += traceStep.getNode().getValue();
        }
        return result;
    }


    static class TreeNode {
        public TreeNode(long v) {
            this.value = v;
        }

        private long value;
        private TreeNode left;
        private TreeNode right;

        public long getValue() {
            return value;
        }

        public void setValue(long value) {
            this.value = value;
        }

        public TreeNode getLeft() {
            return left;
        }

        public void setLeft(TreeNode left) {
            this.left = left;
        }

        public TreeNode getRight() {
            return right;
        }

        public void setRight(TreeNode right) {
            this.right = right;
        }
    }

    static class TraceStep {
        public TraceStep(TraceStep other) {
            this(other.getNode());
            this.leftVisited = other.leftVisited;
            this.rightVisited = other.rightVisited;
        }

        public String toString() {
            return "" + this.getNode().getValue();
        }

        ;
        private TreeNode node;
        private boolean leftVisited;
        private boolean rightVisited;

        public TreeNode getNode() {
            return node;
        }


        public boolean isLeftVisited() {
            return leftVisited;
        }

        public void setLeftVisited(boolean leftVisited) {
            this.leftVisited = leftVisited;
        }

        public boolean isRightVisited() {
            return rightVisited;
        }

        public void setRightVisited(boolean rightVisited) {
            this.rightVisited = rightVisited;
        }

        TraceStep(TreeNode node) {
            this.node = node;
        }
    }


}