[LeetCode] Inorder Successor in BST

Problem Description:

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.

There are just two cases:

1. The easier one: p has right subtree, then its successor is just the leftmost child of its right subtree;
2. The harder one: p has no right subtree, then a traversal is needed to find its successor.

Traversal: we start from the root, each time we see a node with val larger than p -> val, we know this node may be p's successor. So we record it in suc. Then we try to move to the next level of the tree: if p -> val > root -> val, which means p is in the right subtree, then its successor is also in the right subtree, so we update root = root -> right; if p -> val < root -> val, we update root = root -> left similarly; once we find p -> val == root -> val, we know we've reached at p and the current sucis just its successor.

The code is as follows. You may try some examples to see how it works :-)

 1 class Solution {
 2 public:
 3     TreeNode* inorderSuccessor(TreeNode* root, TreeNode* p) {
 4         if (p -> right) return leftMost(p -> right);
 5         TreeNode* suc = NULL;
 6         while (root) {
 7             if (p -> val < root -> val) {
 8                 suc = root;
 9                 root = root -> left;
10             }
11             else if (p -> val > root -> val)
12                 root = root -> right; 
13             else break;
14         }
15         return suc;
16     }
17 private:
18     TreeNode* leftMost(TreeNode* node) {
19         while (node -> left) node = node -> left;
20         return node;
21     }
22 };