【LeetCode 173】Binary Search Tree Iterator

Implement an iterator over a binary search tree (BST). Your iterator will be initialized with the root node of a BST.

Calling next() will return the next smallest number in the BST.

Note: next() and hasNext() should run in average O(1) time and uses O(h) memory, where h is the height of the tree.

 

题意：

　　实现二分搜索树的hasNext() 以及next() 操作，要求 O(1) time and O(h) memory。

思路：

　　暴力的方法是遍历整个树然后将所有元素放入有序的队列中，依次取出，但空间复杂度为O(n)。O(h)的复杂度的解法可通过一个栈来实现：依次将最左边的元素压栈，每次弹出最左下的元素即为最小的元素，同时判断其是否有右子树，若有右子树则继续将其右结点的左边元素依次压栈，循环直到栈为空。

C++：

 1 /**

 2  * Definition for binary tree

 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 BSTIterator {

11 public:

12     

13     BSTIterator(TreeNode *root):_min(0) {

14         while(root != 0)

15         {

16             _stack.push(root);

17             root = root->left;

18         }

19     }

20 

21     /** @return whether we have a next smallest number */

22     bool hasNext() {

23         

24         if(_stack.empty())

25             return false;

26         else

27         {

28             TreeNode* curNode = _stack.top();

29             _min = curNode->val;

30             _stack.pop();

31             

32             if(curNode->right != 0)

33             {

34                 TreeNode* newNode = curNode->right;

35                 while(newNode != 0)

36                 {

37                     _stack.push(newNode);

38                     newNode = newNode->left;

39                 }

40             }

41             return true;

42         }

43     }

44 

45     /** @return the next smallest number */

46     int next() {

47         return _min;

48     }

49 

50 private:

51     stack<TreeNode*> _stack;

52     int _min;

53 };

54 

55 /**

56  * Your BSTIterator will be called like this:

57  * BSTIterator i = BSTIterator(root);

58  * while (i.hasNext()) cout << i.next();

59  */

 

 

Python：

 1 # Definition for a  binary tree node

 2 # class TreeNode:

 3 #     def __init__(self, x):

 4 #         self.val = x

 5 #         self.left = None

 6 #         self.right = None

 7 

 8 class BSTIterator:

 9     # @param root, a binary search tree's root node

10     def __init__(self, root):

11         self.minval = 0

12         self.L = []

13         while root is not None:

14             self.L.append(root)

15             root = root.left

16 

17     # @return a boolean, whether we have a next smallest number

18     def hasNext(self):

19         if len(self.L) == 0:

20             return False

21         else:

22             curNode = self.L.pop()

23             self.minval = curNode.val

24             

25             if curNode.right is not None:

26                 newNode = curNode.right

27                 while newNode is not None:

28                     self.L.append(newNode)

29                     newNode = newNode.left

30             

31             return True

32 

33     # @return an integer, the next smallest number

34     def next(self):

35         return self.minval

36 

37 # Your BSTIterator will be called like this:

38 # i, v = BSTIterator(root), []

39 # while i.hasNext(): v.append(i.next())