173. Binary Search Tree Iterator O(n) time O(h) space
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45
|
class Solution { public: vector<int> getAllElements(TreeNode* root1, TreeNode* root2) { pushAllLeft(root1, s1); pushAllLeft(root2, s2); vector<int> res; while (!s1.empty() || !s2.empty()) { auto p = s1.empty() ? INT_MAX : s1.top()->val; auto q = s2.empty() ? INT_MAX : s2.top()->val; int mn = min(p, q); if (p == mn) { res.push_back(p); pushAllLeft(s1.top()->right, s1); } if (q == mn) { res.push_back(q); pushAllLeft(s2.top()->right, s2); } } return res; }
void pushAllLeft(TreeNode *p, stack<TreeNode *> &s) { if (!s.empty()) { s.pop(); } while (p) { s.push(p); p = p->left; } }
stack<TreeNode *> s1, s2; };
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51
|
class Solution { public: vector<int> getAllElements(TreeNode* root1, TreeNode* root2) { pushAllLeft(root1, s1); pushAllLeft(root2, s2); vector<int> res; while (!s1.empty() || !s2.empty()) { auto p = s1.empty() ? nullptr : s1.top(); auto q = s2.empty() ? nullptr : s2.top(); if (p && q) { if (p->val <= q->val) { res.push_back(p->val); pushAllLeft(p->right, s1); } else { res.push_back(q->val); pushAllLeft(q->right, s2); } } else if (p) { res.push_back(p->val); pushAllLeft(p->right, s1); } else { res.push_back(q->val); pushAllLeft(q->right, s2); } } return res; }
void pushAllLeft(TreeNode *p, stack<TreeNode *> &s) { if (!s.empty()) { s.pop(); } while (p) { s.push(p); p = p->left; } }
stack<TreeNode *> s1, s2; };
|