230. Kth Smallest Element in a BST

My original solution is to traverse the tree in preorder and save the result in an array. Return the k-th number in the array.

1:  class Solution {  
2:  public:  
3:    int kthSmallest(TreeNode* root, int k) {  
4:      vector<int> res;  
5:      helper(root, res);  
6:      return res[k-1];  
7:    }  
8:    void helper(TreeNode* root, vector<int> &res) {  
9:      if (root == NULL) return;  
10:      helper(root->left, res);  
11:      res.push_back(root->val);  
12:      helper(root->right, res);  
13:    }  
14:  };  

Another way is to count the left nodes and use binary search to get the k-th number. However, this is not an optimal solution whose running time is O(NlogN). If we can modify augment the TreeNode data structure and keep track its left child numbers when building the tree, we can achieve the search by O(logN).

1:  /**  
2:   * Definition for a binary tree node.  
3:   * struct TreeNode {  
4:   *   int val;  
5:   *   TreeNode *left;  
6:   *   TreeNode *right;  
7:   *   TreeNode(int x) : val(x), left(NULL), right(NULL) {}  
8:   * };  
9:   */  
10:  class Solution {  
11:  public:  
12:    int kthSmallest(TreeNode* root, int k) {  
13:      int c = countNodes(root->left);  
14:      if (c == k-1) return root->val;   
15:      if (c < k-1) {  
16:        return kthSmallest(root->right, k-c-1);  
17:      } else {  
18:        return kthSmallest(root->left, k);  
19:      }  
20:    }  
21:    int countNodes(TreeNode *root) {  
22:      if (root == NULL) return 0;  
23:      return 1 + countNodes(root->left) + countNodes(root->right);  
24:    }  
25:  };