題目描述
給定一棵非空二叉搜索樹以及一個(gè)target值，找到 BST 中最接近給定值的 k 個(gè)數(shù)。
給出的target值為浮點(diǎn)數(shù)
你可以假設(shè) k 總是合理的，即 k ≤ 總節(jié)點(diǎn)數(shù)
我們可以保證給出的 BST 中只有唯一一個(gè)最接近給定值的 k 個(gè)值的集合
假設(shè)是一棵平衡二叉搜索樹，你可以用時(shí)間復(fù)雜度低于O(n)的算法解決問(wèn)題嗎( n 為節(jié)點(diǎn)個(gè)數(shù))？
測(cè)試樣例
輸入:
{1}
0.000000
1
輸出:
[1]
解釋：
二叉樹 {1}，表示如下的樹結(jié)構(gòu)：
1
輸入:
{3,1,4,#,2}
0.275000
2
輸出:
[1,2]
解釋：
二叉樹 {3,1,4,#,2}，表示如下的樹結(jié)構(gòu)：
? 3
/? \
1? ? 4
\
? 2
解題思路
1、堆棧
設(shè)置兩個(gè)stack，分別從兩個(gè)方向?qū)ふ遗ctarget接近的節(jié)點(diǎn)，其中一個(gè)存儲(chǔ)的值比target大，記為next_stack，另一個(gè)則不大于target，記為prev_stack。
先對(duì)兩個(gè)棧進(jìn)行初始化，當(dāng)target小于當(dāng)前節(jié)點(diǎn)時(shí)，將該節(jié)點(diǎn)入棧next_stack，同時(shí)將節(jié)點(diǎn)置為其左兒子；否則，將該節(jié)點(diǎn)入棧prev_stack，并將節(jié)點(diǎn)置為其右兒子；重復(fù)這個(gè)過(guò)程直至節(jié)點(diǎn)為空。
接下來(lái)進(jìn)入循環(huán)，循環(huán)k次，代表每輪循環(huán)都會(huì)找到一個(gè)接近target的數(shù)。循環(huán)一開始先判斷兩個(gè)stack是否都為空，如果是，則跳出循環(huán)。
然后，比較兩個(gè)stack的棧頂與target的接近程度（若其中一個(gè)棧為空，則將其余target的差值設(shè)置為一個(gè)極大值例如sys.maxsize），若prev_stack的棧頂元素與target較接近，則將其加入到結(jié)果列表中，同時(shí)要在該棧中找到下一個(gè)節(jié)點(diǎn)；否則，對(duì)next_stack進(jìn)行同樣的處理。
最后，待循環(huán)結(jié)束返回找到的k個(gè)數(shù)即可。
另外，這里說(shuō)明下拿出棧頂元素后找到下一個(gè)節(jié)點(diǎn)的方法，以prev_stack為例：
prev_stack中的值是不大于target的，在初始化的時(shí)候。一直是往右子樹的方向中找，從而去逼近target。于是，拿出棧頂元素后，就從棧頂?shù)淖髢鹤娱_始，從這個(gè)左兒子往其右子樹的方向中去找，每找到一個(gè)節(jié)點(diǎn)都將其入棧，并將當(dāng)前節(jié)點(diǎn)置為其右兒子，節(jié)點(diǎn)為空。
為什么要從棧頂?shù)淖髢鹤娱_始而不是從其父節(jié)點(diǎn)開始呢？因?yàn)闂ｍ斒瞧涓腹?jié)點(diǎn)的右兒子，因此棧頂?shù)淖髢鹤与m比棧頂小但比其父節(jié)點(diǎn)還是大的，相比較而言更接近target，因?yàn)閜rev_stack中都是小于target的，所以在prev_stack中，越大就越接近target。
2、先序遍歷
從根節(jié)點(diǎn)開始往左子樹方向遍歷，每遇到一個(gè)節(jié)點(diǎn)就加入棧中，同時(shí)將節(jié)點(diǎn)更新為左兒子，重復(fù)這個(gè)過(guò)程直至節(jié)點(diǎn)為空。
然后，進(jìn)入循環(huán)，直至棧空。每次循環(huán)都彈出棧頂，將其加入到結(jié)果列表中，使得列表是從小到大排序的。
循環(huán)一開始，需要判斷下列表中是否已經(jīng)有超過(guò)k個(gè)數(shù)，如果是并且當(dāng)前棧頂?shù)倪@個(gè)數(shù)與target的接近程度不如列表中的倒數(shù)第k個(gè)數(shù)，那么就可以跳出循環(huán)了，因?yàn)闂：罄m(xù)的元素都比棧頂大，那么它們與target的接近程度就更加遠(yuǎn)了。
當(dāng)棧頂加入到結(jié)果列表后，將其更新為右兒子，若右兒子不為空，則入棧，并從這個(gè)右兒子的左子樹方向一直遞歸，每到一個(gè)節(jié)點(diǎn)就將其入棧，直至節(jié)點(diǎn)為空。
以上就是每次循環(huán)的操作，循環(huán)結(jié)束后返回結(jié)果列表的倒數(shù)前k個(gè)元素即為所求。
代碼
1、堆棧
import sys
"""
Definition of TreeNode:
class TreeNode:
? ? def __init__(self, val):
? ? ? ? self.val = val
? ? ? ? self.left, self.right = None, None
"""
class Solution:
? ? """
? ? @param root: the given BST
? ? @param target: the given target
? ? @param k: the given k
? ? @return: k values in the BST that are closest to the target
? ? """
? ? def closestKValues(self, root, target, k):
? ? ? ? if not root or not k:
? ? ? ? ? ? return []
? ? ? ? closest_ks = []
? ? ? ? prev_stack, next_stack = self.bulid_stacks(root, target)
? ? ? ? for _ in range(k):
? ? ? ? ? ? if not (len(prev_stack) or len(next_stack)):
? ? ? ? ? ? ? ? break
? ? ? ? ? ? delta_prev = abs(prev_stack[-1].val - target) if prev_stack else sys.maxsize
? ? ? ? ? ? delta_next = abs(next_stack[-1].val - target) if next_stack else sys.maxsize
? ? ? ? ? ? if delta_prev < delta_next:
? ? ? ? ? ? ? ? closest_ks.append(self.get_prev(prev_stack))
? ? ? ? ? ? else:
? ? ? ? ? ? ? ? closest_ks.append(self.get_next(next_stack))
? ? ? ? return closest_ks
? ? def bulid_stacks(self, node, target):
? ? ? ? prev_stack, next_stack = [], []
? ? ? ? # 從兩個(gè)方向逼近target
? ? ? ? while node:
? ? ? ? ? ? # next_stack往值小的方向走，但值比target大
? ? ? ? ? ? if target < node.val:
? ? ? ? ? ? ? ? next_stack.append(node)
? ? ? ? ? ? ? ? node = node.left
? ? ? ? ? ? # prev_stack往值大的方向走，但值比target小
? ? ? ? ? ? else:
? ? ? ? ? ? ? ? prev_stack.append(node)
? ? ? ? ? ? ? ? node = node.right
? ? ? ? return prev_stack, next_stack
? ? def get_prev(self, prev_stack):
? ? ? ? node = prev_stack.pop()
? ? ? ? value = node.val
? ? ? ? # prev_stack中，上次已經(jīng)入棧了所有右子樹，
? ? ? ? # 因此彈出當(dāng)前節(jié)點(diǎn)后，要繼續(xù)逼近target，
? ? ? ? # 就要從其左兒子的右子樹開始
? ? ? ? node = node.left
? ? ? ? while node:
? ? ? ? ? ? prev_stack.append(node)
? ? ? ? ? ? node = node.right
? ? ? ? return value
? ? def get_next(self, next_stack):
? ? ? ? node = next_stack.pop()
? ? ? ? value = node.val
? ? ? ? # next_stack中，上次已經(jīng)入棧了所有左子樹
? ? ? ? # 因此彈出當(dāng)前節(jié)點(diǎn)后，要繼續(xù)逼近target，
? ? ? ? # 就要從其右兒子的左子樹開始
? ? ? ? node = node.right
? ? ? ? while node:
? ? ? ? ? ? next_stack.append(node)
? ? ? ? ? ? node = node.left
? ? ? ? return value
2、先序遍歷
"""
Definition of TreeNode:
class TreeNode:
? ? def __init__(self, val):
? ? ? ? self.val = val
? ? ? ? self.left, self.right = None, None
"""
class Solution:
? ? """
? ? @param root: the given BST
? ? @param target: the given target
? ? @param k: the given k
? ? @return: k values in the BST that are closest to the target
? ? """
? ? def closestKValues(self, root, target, k):
? ? ? ? ans = []
? ? ? ? stack = []
? ? ? ? # 先序遍歷，從左路一直下去，棧頂獲得最小值
? ? ? ? while root:
? ? ? ? ? ? stack.append(root)
? ? ? ? ? ? root = root.left
? ? ? ? while stack:
? ? ? ? ? ? node = stack.pop()
? ? ? ? ? ? # ans中的節(jié)點(diǎn)是根據(jù)值從小到大排序的
? ? ? ? ? ? # 若當(dāng)前其中已經(jīng)超過(guò)了k個(gè)，并且倒數(shù)第k個(gè)與target的差
? ? ? ? ? ? # 已經(jīng)小于當(dāng)前棧彈出的節(jié)點(diǎn)，那么就已經(jīng)獲得了最接近target的k個(gè)數(shù)了
? ? ? ? ? ? # 因?yàn)闂Ｖ蟮墓?jié)點(diǎn)與target的差都比棧頂?shù)倪@個(gè)節(jié)點(diǎn)大
? ? ? ? ? ? if len(ans) >= k and abs(ans[-k] - target) < abs(node.val - target):
? ? ? ? ? ? ? ? break
? ? ? ? ? ? ans.append(node.val)
? ? ? ? ? ? node = node.right
? ? ? ? ? ? while node:
? ? ? ? ? ? ? ? stack.append(node)
? ? ? ? ? ? ? ? node = node.left
? ? ? ? # 最后返回倒數(shù)前k嘅數(shù)即為所求
? ? ? ? return ans[-k:]
