Friday, October 9, 2015

Leetcode: 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.
Brute-force solution: O(n)
The most straight-forward solution is to do a in-order traversal of the BST. When we found the target node, we look for the smallest number greater than the node.

Code (Java):
/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
public class Solution {
    public TreeNode inorderSuccessor(TreeNode root, TreeNode p) {
        if (root == null || p == null) {
            return null;
        }
        
        if (p == root) {
            return p.right;
        }
        
        Stack<TreeNode> stack = new Stack<>();
        TreeNode q = root;
        
        while (!stack.isEmpty() || q != null) {
            if (q != null) {
                stack.push(q);
                q = q.left;
            } else {
                TreeNode curr = stack.pop();
                q = curr.right;
                
                if (curr == p) {
                    if (curr.right != null) {
                        TreeNode m = curr.right;
                        while (m.left != null) {
                            m = m.left;
                        }
                        return m;
                        
                    } else if (!stack.isEmpty()){
                        return stack.pop();
                    }
                }
            }
        }
        
        return null;
    }
}

An O(h) Solution:
Since the given p node is a node in the BST (not just a value), we can directly start from the node.  There are two cases to consider:
1. If the p.right != null. Then the successor of p must be the minimum number of the right subtree.
2. If the p.right == null, then the successor of p must be the minimum number which is greater than p's value. We could start from root, if the value of the root is greater than the p.val. Then we go to the left subtree because we wanna try a smaller one. However we need to save this node because it is possible that the root.left would be smaller than the p.val. If the root's val is less than the p.val, then we simply go to the right subtree. 

Code (Java):
/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
public class Solution {
    public TreeNode inorderSuccessor(TreeNode root, TreeNode p) {
        if (root == null || p == null) {
            return null;
        }
        
        if (p.right != null) {
            return findMin(p.right);
        }
        
        // Case 2: p.right == null
        TreeNode succ = null;
        TreeNode q = root;
        
        while (q != null) {
            if (q.val > p.val) {
                succ = q;
                q = q.left;
            } else if (q.val < p.val) {
                q = q.right;
            } else {
                break;
            }
        }
        
        return succ;
    }
    
    private TreeNode findMin(TreeNode root) {
        TreeNode p = root;
        
        while (p.left != null) {
            p = p.left;
        }
        
        return p;
    }
}

2 comments: