LeetCode Offer 54. kth largest node in BST

class Solution:  # recursion
  def kthLargest(self, root: TreeNode, k: int) -> int:
    def dfs(node):
      nonlocal k
      if node.right and (right := dfs(node.right)) != None:
        return right
      k -= 1
      if k == 0:
        return node.val
      if node.left and (left := dfs(node.left)) != None:
        return left
    return dfs(root)

class Solution:
  def kthLargest(self, root: TreeNode, k: int) -> int:
    count, ans = 0, None

    def dfs(root):
      nonlocal count, ans
      if root and ans == None:
        dfs(root.right)
        count += 1
        if count == k:
          ans = root.val
        dfs(root.left)

    dfs(root)
    return ans

class Solution:
  def kthLargest(self, root: TreeNode, k: int) -> int:
    # 第 k 大的节点，比如第 1 大的就是最大的那个
    # 所以是中序遍历，且从右到左
    # 迭代写法（最好记忆的颜色标记法）
    stack = [(root, False)]  # 节点, 是否取出
    count = 0
    while stack:
      node, pick = stack.pop()
      if pick:
        count += 1
        if count == k:
          return node.val
      else:
        if node.left: stack.append((node.left, False))
        stack.append((node, True))
        if node.right: stack.append((node.right, False))