Leetcode 230: Kth Smallest Element in a BST

Posted 2020-10-16 Keep walking

Given a binary search tree, write a function kthSmallest to find the kth smallest element in it.

Note: 
You may assume k is always valid, 1 ≤ k ≤ BST‘s total elements.

Follow up:
What if the BST is modified (insert/delete operations) often and you need to find the kth smallest frequently? How would you optimize the kthSmallest routine?

 

 1 /**
 2  * Definition for a binary tree node.
 3  * public class TreeNode {
 4  *     public int val;
 5  *     public TreeNode left;
 6  *     public TreeNode right;
 7  *     public TreeNode(int x) { val = x; }
 8  * }
 9  */
10 public class Solution {
11     private int CountNodes(TreeNode node)
12     {
13         if (node == null) return 0;
14         return 1 + CountNodes(node.left) + CountNodes(node.right);
15     }
16     
17     public int KthSmallest(TreeNode root, int k) {
18         
19         int left = CountNodes(root.left);
20         
21         if (k <= left) 
22         {
23             return KthSmallest(root.left, k);
24         }
25         else if (k > left + 1)
26         {        
27             return KthSmallest(root.right, k - left - 1);
28         }
29         
30         return root.val;
31     }
32 }
33 
34 public class Solution1 {
35     public int KthSmallest(TreeNode root, int k) {
36         var stack = new Stack<TreeNode>();
37         
38         while (stack.Count > 0 || root != null)
39         {
40             if (root != null)
41             {
42                 stack.Push(root);
43                 root = root.left;
44             }
45             else
46             {
47                 root = stack.Pop();
48                 if (--k == 0) return root.val;
49                 
50                 // after pop, we must set the current node the right node otherwise we will see endless loop
51                 root = root.right;               
52             }
53         }
54         
55         return 0;
56     }
57 }
58 
59 public class Solution2 {
60     public int KthSmallest(TreeNode root, int k) {
61         var result = new int[] {0, 0};
62         
63         InOrder(root, k, result);
64         
65         return result[1];
66     }
67     
68     private void InOrder(TreeNode node, int k, int[] result)
69     {
70         if (node == null) return;
71         
72         InOrder(node.left, k, result);
73         
74         result[0]++;
75         
76         if (result[0] == k)
77         {
78             result[1] = node.val;
79             return;
80         }
81         
82         InOrder(node.right, k, result);
83     }
84 }