Question

面试题 54. 二叉搜索树的第 k 大节点

题目描述

给定一棵二叉搜索树，请找出其中第 k 大的节点的值。

 

示例 1:

输入: root = [3,1,4,null,2], k = 1
   3
  / \
 1   4
  \
   2
输出: 4

示例 2:

输入: root = [5,3,6,2,4,null,null,1], k = 3
     5
    / \
   3   6
  / \
   2   4
  /
 1
输出: 4

 

限制：

• 1 ≤ k ≤ 二叉搜索树元素个数

解法

方法一：反序中序遍历

由于二叉搜索树的中序遍历是升序的，因此可以反序中序遍历，即先递归遍历右子树，再访问根节点，最后递归遍历左子树。

这样就可以得到一个降序的序列，第 $k$ 个节点就是第 $k$ 大的节点。

时间复杂度 $O(n)$，空间复杂度 $O(n)$。其中 $n$ 是二叉搜索树的节点个数。

# Definition for a binary tree node.
# class TreeNode:
#   def __init__(self, x):
#     self.val = x
#     self.left = None
#     self.right = None


class Solution:
  def kthLargest(self, root: TreeNode, k: int) -> int:
    def dfs(root):
      nonlocal k, ans
      if root is None or k == 0:
        return
      dfs(root.right)
      k -= 1
      if k == 0:
        ans = root.val
      dfs(root.left)

    ans = 0
    dfs(root)
    return ans

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *   int val;
 *   TreeNode left;
 *   TreeNode right;
 *   TreeNode(int x) { val = x; }
 * }
 */
class Solution {
  private int k;
  private int ans;

  public int kthLargest(TreeNode root, int k) {
    this.k = k;
    dfs(root);
    return ans;
  }

  private void dfs(TreeNode root) {
    if (root == null || k == 0) {
      return;
    }
    dfs(root.right);
    if (--k == 0) {
      ans = root.val;
    }
    dfs(root.left);
  }
}

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *   int val;
 *   TreeNode *left;
 *   TreeNode *right;
 *   TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
  int kthLargest(TreeNode* root, int k) {
    int ans = 0;
    function<void(TreeNode*)> dfs = [&](TreeNode* root) {
      if (!root || !k) {
        return;
      }
      dfs(root->right);
      if (--k == 0) {
        ans = root->val;
      }
      dfs(root->left);
    };
    dfs(root);
    return ans;
  }
};

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *   Val int
 *   Left *TreeNode
 *   Right *TreeNode
 * }
 */
func kthLargest(root *TreeNode, k int) (ans int) {
  var dfs func(*TreeNode)
  dfs = func(root *TreeNode) {
    if root == nil || k == 0 {
      return
    }
    dfs(root.Right)
    k--
    if k == 0 {
      ans = root.Val
    }
    dfs(root.Left)
  }
  dfs(root)
  return
}

/**
 * Definition for a binary tree node.
 * class TreeNode {
 *   val: number
 *   left: TreeNode | null
 *   right: TreeNode | null
 *   constructor(val?: number, left?: TreeNode | null, right?: TreeNode | null) {
 *     this.val = (val===undefined ? 0 : val)
 *     this.left = (left===undefined ? null : left)
 *     this.right = (right===undefined ? null : right)
 *   }
 * }
 */

function kthLargest(root: TreeNode | null, k: number): number {
  let res = 0;
  const dfs = (root: TreeNode | null) => {
    if (root == null) {
      return;
    }
    dfs(root.right);
    if (--k === 0) {
      res = root.val;
      return;
    }
    dfs(root.left);
  };
  dfs(root);
  return res;
}

// Definition for a binary tree node.
// #[derive(Debug, PartialEq, Eq)]
// pub struct TreeNode {
//   pub val: i32,
//   pub left: Option<Rc<RefCell<TreeNode>>>,
//   pub right: Option<Rc<RefCell<TreeNode>>>,
// }
//
// impl TreeNode {
//   #[inline]
//   pub fn new(val: i32) -> Self {
//   TreeNode {
//     val,
//     left: None,
//     right: None
//   }
//   }
// }
use std::rc::Rc;
use std::cell::RefCell;

impl Solution {
  fn dfs(root: &Option<Rc<RefCell<TreeNode>>>, arr: &mut Vec<i32>) {
    if let Some(node) = root {
      let node = node.borrow();
      Solution::dfs(&node.right, arr);
      arr.push(node.val);
      Solution::dfs(&node.left, arr)
    }
  }
  pub fn kth_largest(root: Option<Rc<RefCell<TreeNode>>>, k: i32) -> i32 {
    let mut arr = vec![];
    Solution::dfs(&root, &mut arr);
    arr[(k - 1) as usize]
  }
}

/**
 * Definition for a binary tree node.
 * function TreeNode(val) {
 *   this.val = val;
 *   this.left = this.right = null;
 * }
 */
/**
 * @param {TreeNode} root
 * @param {number} k
 * @return {number}
 */
var kthLargest = function (root, k) {
  let ans = 0;
  const dfs = root => {
    if (!root || !k) {
      return;
    }
    dfs(root.right);
    if (--k == 0) {
      ans = root.val;
    }
    dfs(root.left);
  };
  dfs(root);
  return ans;
};

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *   public int val;
 *   public TreeNode left;
 *   public TreeNode right;
 *   public TreeNode(int x) { val = x; }
 * }
 */
public class Solution {
  private int ans;
  private int k;

  public int KthLargest(TreeNode root, int k) {
    this.k = k;
    dfs(root);
    return ans;
  }

  private void dfs(TreeNode root) {
    if (root == null || k == 0) {
      return;
    }
    dfs(root.right);
    if (--k == 0) {
      ans = root.val;
    }
    dfs(root.left);
  }
}

方法二

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *   Val int
 *   Left *TreeNode
 *   Right *TreeNode
 * }
 */
func kthLargest(root *TreeNode, k int) int {
  ch := make(chan int)
  ctx, cancel := context.WithCancel(context.Background())
  defer cancel()
  go inorder(ctx, root, ch)
  for ; k > 1; k-- {
    <-ch
  }
  return <-ch
}

func inorder(ctx context.Context, cur *TreeNode, ch chan<- int) {
  if cur != nil {
    inorder(ctx, cur.Right, ch)
    select {
    case ch <- cur.Val:
    case <-ctx.Done():
      return
    }
    inorder(ctx, cur.Left, ch)
  }
}