CareerCup 4.6

4.6 Write an algorithm to find the 'next' node (i.e., in-order successor) of a given node in a binary search tree. You may assume that each node has a link to its parent.

struct TreeNode {
    int val;
    TreeNode *parent;
    TreeNode *left;
    TreeNode *right;
    TreeNode(int x) : val(x), parent(NULL), left(NULL), right(NULL) {}
};

TreeNode* leftMostNode(TreeNode *node) {
    if (node == NULL) {
        return NULL;
    }
    while (node->left) {
        node = node->left;
    }
    return node;
}

TreeNode* inorderSuccesor(TreeNode *node) {
    if (node == NULL) {
        return NULL;
    }

    if (node->parent == NULL || node->right) {
        return leftMostNode(node->right);
    } else {
        TreeNode *p = node;
        TreeNode *q = p->parent;
        while (q && q->left != p) {
            p = q;
            q = q->parent;
        }
        return q;
    }
}