173. 二叉搜索树迭代器
分别结合 栈 和 队列 两种数据结构实现。
class BSTIterator:
"""
结合栈
173. 二叉搜索树迭代器
https://leetcode.cn/problems/binary-search-tree-iterator/description/?show=1
"""
def __init__(self, root: Optional[TreeNode]):
self.stk = []
self.pushleftBranch(root)
# 左侧树枝一撸到底
def pushleftBranch(self, p):
while p:
self.stk.append(p)
p = p.left
def next(self) -> int:
p = self.stk.pop()
self.pushleftBranch(p.right)
return p.val
def hasNext(self) -> bool:
return len(self.stk) > 0
def peek(self) -> bool:
return self.stk[-1].val
class BSTIterator1:
"""
结合队列
173. 二叉搜索树迭代器
https://leetcode.cn/problems/binary-search-tree-iterator/description/?show=1
"""
def __init__(self, root: Optional[TreeNode]):
self.deque = collections.deque()
self.inOrder(root)
def inOrder(self, root):
"""
二叉树中序遍历
:param root:
:return:
"""
if not root:
return
self.inOrder(root.left)
self.deque.append(root.value)
self.inOrder(root.right)
def next(self) -> int:
val = self.deque.popleft()
return val
def hasNext(self) -> bool:
return len(self.deque) > 0
def peek(self) -> bool:
return self.deque[0]
BST迭代器应用
# Definition for a binary tree node.
# class TreeNode:
# def __init__(self, val=0, left=None, right=None):
# self.val = val
# self.left = left
# self.right = right
class Solution:
"""
1305. 两棵二叉搜索树中的所有元素
"""
def getAllElements(self, root1: TreeNode, root2: TreeNode) -> List[int]:
t1 = BSTIterator(root1)
t2 = BSTIterator(root2)
res = []
while t1.hasNext() and t2.hasNext():
if t1.peek() > t2.peek():
res.append(t2.next())
else:
res.append(t1.next())
while t1.hasNext():
res.append(t1.next())
while t2.hasNext():
res.append(t2.next())
return res
class BSTIterator:
"""
结合队列
173. 二叉搜索树迭代器
https://leetcode.cn/problems/binary-search-tree-iterator/description/?show=1
"""
def __init__(self, root: Optional[TreeNode]):
self.deque = collections.deque()
self.inOrder(root)
def inOrder(self, root):
"""
二叉树中序遍历
:param root:
:return:
"""
if not root:
return
self.inOrder(root.left)
self.deque.append(root.val)
self.inOrder(root.right)
def next(self) -> int:
val = self.deque.popleft()
return val
def hasNext(self) -> bool:
return len(self.deque) > 0
def peek(self) -> bool:
return self.deque[0]