Inorder Successor in BST

题目
Given a binary search tree and a node in it, find the in-order successor of that node in the BST.

Note: If the given node has no in-order successor in the tree, return null.

答案

class Solution {
    // Iterate the tree, if p is found, then next treenode to be visited is its successor
    public TreeNode inorderSuccessor(TreeNode root, TreeNode p) {
        if(root == null) return null;
        TreeNode ans = null;
        while(root != null) {
            if(root.val > p.val) {
                // root could be p's successor 
                ans = root;
                root = root.left;
            }
            else {
                root = root.right;
            }
        }
        return ans;
    }
}