Inorder Successor in BST Python Solution

lcpy gen -n 285   # by problem number
lcpy gen -s inorder_successor_in_bst   # by problem name

Input: root = [2,1,3], p = 1
Output: 2
Explanation: 1's in-order successor node is 2. Note that both p and the return value are of TreeNode type.

Input: root = [5,3,6,2,4,null,null,1], p = 6
Output: null
Explanation: There is no in-order successor of the current node, so the answer is null.

from leetcode_py import TreeNode


class Solution:
    # Time: O(h)
    # Space: O(1)
    def inorder_successor(
        self, root: TreeNode[int] | None, p: TreeNode[int]
    ) -> TreeNode[int] | None:
        # If p has a right subtree, successor is leftmost node of that subtree.
        # Otherwise, successor is the lowest ancestor for which p lies in its
        # left subtree. Track candidate successor while walking down.
        successor = None
        while root:
            if p.val < root.val:
                successor = root
                root = root.left
            else:
                root = root.right
        return successor