LeetCode 230 Kth Smallest Element in a BST

LeetCode

Given the root of a binary search tree, and an integer k, return the k<sup>th</sup> smallest value (1-indexed) of all the values of the nodes in the tree.

Example 1:

Input: root = [3,1,4,null,2], k = 1
Output: 1

Example 2:

Input: root = [5,3,6,2,4,null,null,1], k = 3
Output: 3

Constraints:

The number of nodes in the tree is n. 1 <= k <= n <= 10<sup>4</sup>
0 <= Node.val <= 10<sup>4</sup>

*Follow up: If the BST is modified often (i.e., we can do insert and delete operations) and you need to find the kth smallest frequently, how would you optimize?

题目大意：

求BST的第k个最小元素（k从1开始）

算法思路：

N/A

注意事项：

1. 用返回值(count, 结果). 递归右节点，用k - 1 - left_count
2. Line 9不能用if left_val，因为left_val可能是0，用is not None

Python代码：

def kthSmallest(self, root: TreeNode, k: int) -> int:
	count, val = self.dfs(root, k)
	return val

def dfs(self, root, k):
	if not root:
		return 0, None

	left_count, left_val = self.dfs(root.left, k)
	if left_val is not None: # remember if left_val is wrong
		return 0, left_val
	if left_count + 1 == k:
		return 0, root.val
	right_count, right_val = self.dfs(root.right, k - left_count - 1) # remember k-1-left_count
    if right_val is not None:
		return 0, right_val
	return left_count + 1 + right_count, right_val

Java代码：

public int kthSmallest2(TreeNode root, int k) {
	return kthSmallest2R(root, k).result;
}

public ResultType kthSmallest2R(TreeNode root, int k) {
	if (root == null)
		return new ResultType(null, 0);
	ResultType leftResult = kthSmallest2R(root.left, k);

	Integer result = leftResult.result;
	if(result != null)
		return new ResultType(result, 0);
	if(leftResult.count + 1 == k)
		return new ResultType(root.val, 0);
	
	ResultType rightResult = kthSmallest2R(root.right, k - 1 - leftResult.count);
	result = rightResult.result;
	
	return new ResultType(result, leftResult.count + 1 + rightResult.count);
}

class ResultType {
	Integer result;
	int count;
	
	public ResultType(Integer r, int c) {
		result = r;
		count = c;
	}
}

算法分析：

时间复杂度为O(n)，空间复杂度O(1)