Inorder Successor in Binary Search Tree

Given a binary search tree (See Definition) and a node in it, find the in-order successor of that node in the BST.

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

Solution:

in-order traversal is from left subtree to current node then right subtree. This gives us following rule:

1. if the given node is on left subtree, the current node is successor when the given node has no right child.
2. if the given node has right child, the most left child on right subtree is the successor

Code:

public class Solution {
    public TreeNode inorderSuccessor(TreeNode root, TreeNode p) {
        if (root == null || p == null) {
            return null;
        }
        TreeNode cur = p.right;
        if (cur != null) {
            while  (cur.left != null) {
                cur = cur.left;
            }
            return cur;
        }
        while (root != null && root != p) {
            if (root.val > p.val) {
                cur = root;
                root = root.left;
            } else {
                root = root.right;
            }
        }
        return cur;
    }
}