https://leetcode.com/problems/closest-binary-search-tree-value-ii/
Closest Binary Search Tree Value II
Given a non-empty binary search tree and a target value, find k values in the BST that are closest to the target.
Note:
- Given target value is a floating point.
- You may assume k is always valid, that is: k ≤ total nodes.
- You are guaranteed to have only one unique set of k values in the BST that are closest to the target.
Follow up:
Assume that the BST is balanced, could you solve it in less than O(n) runtime (where n = total nodes)?
Code
# Definition for a binary tree node.
# class TreeNode(object):
# def __init__(self, x):
# self.val = x
# self.left = None
# self.right = None
class Solution(object):
def closestKValues(self, root, target, k):
"""
:type root: TreeNode
:type target: float
:type k: int
:rtype: List[int]
"""
stk1 = []
stk2 = []
self.inorder(root, False, target, stk1)
self.inorder(root, True, target, stk2)
res = []
for _ in xrange(k):
if not stk1:
res.append(stk2.pop())
elif not stk2:
res.append(stk1.pop())
else:
if abs(stk1[len(stk1)-1] - target) < abs(stk2[len(stk2)-1] - target):
res.append(stk1.pop())
else:
res.append(stk2.pop())
return res
def inorder(self, root, reverse, target, stk):
if root is None:
return
self.inorder(root.right, reverse, target, stk) if reverse else self.inorder(root.left, reverse, target, stk)
# The first inequality is less than or equal, the second inequality must be larger than (without equal).
# Or the first is less than, the second is larger than or equal to
if not reverse and target <= root.val:
return
if reverse and target > root.val:
return
stk.append(root.val)
self.inorder(root.left, reverse, target, stk) if reverse else self.inorder(root.right, reverse, target, stk)
Idea:
If we do inorder depth-first traverse (left->root->right) of a binary search tree, we can get an ascending sorted list of elements. If we do reverse inorder depth-first traverse (right->root->left) of a binary search tree, we can get a descending sorted list.
So we create two stacks to store. Stack 1 records the ascending sorted list but it terminates before the element which <= target. (line 40). Stack 2 records the descending sorted list but it terminates at the element which < target. (line 42). (You can also make the first inequality with less than, and the second inequality with larger than or equal to. The main point is that the two stacks shouldn’t contain any element at the same time.)
At last you peek at stack 1 and stack 2. You need to pop element into `res` list from whichever stack has the top element closer to target.
Reference:
https://leetcode.com/discuss/55240/ac-clean-java-solution-using-two-stacks