Inorder Successor in BST

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.

It's guaranteed p is one node in the given tree. (You can directly compare the memory address to find p)

Example

Example 1:

Input: {1,#,2}, node with value 1
Output: 2
Explanation:
  1
   \
    2

Example 2:

Input: {2,1,3}, node with value 1
Output: 2
Explanation: 
    2
   / \
  1   3

Binary Tree Representation

Challenge

O(h), where h is the height of the BST.

public class Solution {
    TreeNode lastLeft;
    TreeNode ans;

    public TreeNode inorderSuccessor(TreeNode root, TreeNode p) {
        helper(root, p);
        return ans;
    }

    public void helper(TreeNode root, TreeNode target) {
        if (root == null)
            return;
        if (root == target) {
            // Case 1 -> if right subtree exists ,
            // then ans is smallest node from that subtree
            if (root.right != null) {
                TreeNode temp = root.right;
                while (temp.left != null)
                    temp = temp.left;
                ans = temp;
            }
            // Else the answer is the node where we took the last left turn
            // to reach target node
            else
                ans = lastLeft;
        } else {
            if (target.val > root.val)
                helper(root.right, target);
            else {
                lastLeft = root;
                helper(root.left, target);
            }
        }
    }
}