1: class Solution {
2: public:
3: vector<int> closestKValues(TreeNode* root, double target, int k) {
4: vector<int> nodes;
5: vector<int> res;
6: inorder(root, nodes);
7: if (nodes.size() == 0) return res;
8: int l = 0, r = nodes.size()-1;
9: if (target < nodes[l]) {
10: while (k--) res.push_back(nodes[l++]);
11: return res;
12: }
13: if (target > nodes[r]) {
14: while (k--) res.push_back(nodes[r--]);
15: return res;
16: }
17: while (l <= r) {
18: int mid = l + (r - l) / 2;
19: if (nodes[mid] == target) {r = mid; break;}
20: else if (nodes[mid] > target) r = mid - 1;
21: else l = mid+1;
22: }
23: l = r;
24: r = l+1;
25: while (k--) {
26: if (l < 0) res.push_back(nodes[r++]);
27: else if (r == nodes.size()) res.push_back(nodes[l--]);
28: else if (abs(nodes[l]-target) < abs(nodes[r]-target)) {
29: res.push_back(nodes[l--]);
30: } else {
31: res.push_back(nodes[r++]);
32: }
33: }
34: return res;
35: }
36: void inorder(TreeNode *root, vector<int> &nodes) {
37: if (root == NULL) return;
38: inorder(root->left, nodes);
39: nodes.push_back(root->val);
40: inorder(root->right, nodes);
41: }
42: };
There is another way to maintain the predecessor and successor stack by traversing the tree in inorder and reverse-inorder.
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: vector<int> closestKValues(TreeNode* root, double target, int k) {
13: stack<int> predecessor;
14: stack<int> successor;
15: inorder(root, false, predecessor, target);
16: inorder(root, true, successor, target);
17: vector<int> res;
18: while (k--) {
19: if (predecessor.empty()) {
20: res.push_back(successor.top());
21: successor.pop();
22: } else if (successor.empty()) {
23: res.push_back(predecessor.top());
24: predecessor.pop();
25: } else if (abs(predecessor.top()-target) < abs(successor.top()-target)) {
26: res.push_back(predecessor.top());
27: predecessor.pop();
28: } else {
29: res.push_back(successor.top());
30: successor.pop();
31: }
32: }
33: return res;
34: }
35: void inorder(TreeNode *root, bool reverse, stack<int> &stk, double target) {
36: if (root == NULL) return;
37: inorder(reverse ? root->right : root->left, reverse, stk, target);
38: if ((reverse && root->val <= target) || ((!reverse) && root->val > target)) return;
39: stk.push(root->val);
40: inorder(reverse ? root->left : root->right, reverse, stk, target);
41: }
42: };
And this problem actually can be converted to a design problem with two methods getPredecessor() and getSuccessor(). We can stop searching the BST when we find the closest predecessor and successor to target. And we can update the predecessor stack and successor stack in these two methods respectively. So the running time will be O(klogn);
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: private:
12: stack<treenode> pred;
13: stack<treenode> succ;
14: public:
15: vector<int> closestKValues(TreeNode* root, double target, int k) {
16: vector<int> res;
17: initPredecessor(root, target);
18: initSuccessor(root, target);
19: if (!succ.empty() && !pred.empty() && succ.top()->val == pred.top()->val) {
20: getNextPredecessor();
21: }
22: while (k--) {
23: if (succ.empty()) res.push_back(getNextPredecessor());
24: else if (pred.empty()) res.push_back(getNextSuccessor());
25: else if (abs(succ.top()->val - target) < abs(pred.top()->val - target)) {
26: res.push_back(getNextSuccessor());
27: } else {
28: res.push_back(getNextPredecessor());
29: }
30: }
31: return res;
32: }
33: void initPredecessor(TreeNode *root, double target) {
34: while (root) {
35: if (root->val == target) {
36: pred.push(root);
37: break;
38: } else if (root->val < target) {
39: pred.push(root);
40: root = root->right;
41: } else {
42: root = root->left;
43: }
44: }
45: }
46: void initSuccessor(TreeNode *root, double target) {
47: while (root) {
48: if (root->val == target) {
49: succ.push(root);
50: break;
51: } else if (root->val > target) {
52: succ.push(root);
53: root = root->left;
54: } else {
55: root = root->right;
56: }
57: }
58: }
59: int getNextPredecessor() {
60: TreeNode *root = pred.top();
61: pred.pop();
62: int res = root->val;
63: root = root->left;
64: while (root) {
65: pred.push(root);
66: root = root->right;
67: }
68: return res;
69: }
70: int getNextSuccessor() {
71: TreeNode *root = succ.top();
72: succ.pop();
73: int res = root->val;
74: root = root->right;
75: while (root) {
76: succ.push(root);
77: root = root->left;
78: }
79: return res;
80: }
81: };
The second time I revisited this problem, I found that initPredecessor() and initSuccessor() can be combined into one function.
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: private:
12: stack<TreeNode *> pred;
13: stack<TreeNode *> succ;
14: public:
15: vector<int> closestKValues(TreeNode* root, double target, int k) {
16: vector<int> res;
17: if (root == NULL) return res;
18: initStacks(root, target);
19: while (k) {
20: if (pred.empty()) res.push_back(getNextPredecessor());
21: else if (succ.empty()) res.push_back(getNextSuccessor());
22: else if (abs(pred.top()->val-target) < abs(succ.top()->val-target)) {
23: res.push_back(getNextPredecessor());
24: } else {
25: res.push_back(getNextSuccessor());
26: }
27: k--;
28: }
29: return res;
30: }
31: void initStacks(TreeNode *root, double target) {
32: while (root != NULL) {
33: if (root->val <= target) {
34: pred.push(root);
35: root = root->right;
36: } else {
37: succ.push(root);
38: root = root->left;
39: }
40: }
41: }
42: int getNextPredecessor() {
43: TreeNode *t = pred.top();
44: int ret = t->val;
45: pred.pop();
46: t = t->left;
47: while (t) {
48: pred.push(t);
49: t = t->right;
50: }
51: return ret;
52: }
53: int getNextSuccessor() {
54: TreeNode *t = succ.top();
55: int ret = t->val;
56: succ.pop();
57: t = t->right;
58: while (t) {
59: succ.push(t);
60: t = t->left;
61: }
62: return ret;
63: }
64: };
No comments:
Post a Comment