dfs 二叉树中序遍历迭代解法——求解BST中第k小元素

BST中第K小的元素

中文English

给一棵二叉搜索树，写一个 KthSmallest 函数来找到其中第 K 小的元素。

Example

样例 1:

输入：{1,#,2},2
输出：2
解释：
	1
	 \
	  2
第二小的元素是2。

样例 2:

输入：{2,1,3},1
输出：1
解释：
  2
 / \
1   3
第一小的元素是1。

Challenge

如果这棵 BST 经常会被修改(插入/删除操作)并且你需要很快速的找到第 K 小的元素呢？你会如何优化这个 KthSmallest 函数？

Notice

你可以假设 k 总是有效的， 1 ≤ k ≤ 树的总节点数。

"""
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 k: the given k
    @return: the kth smallest element in BST
    """

    """
    nth = 0
    result = None

    def kthSmallest(self, root, k):
        # write your code here
        self.dfs(root, k)
        return self.result

    def dfs(self, root, k):
        if not root:
            return
        self.dfs(root.left, k)
        self.nth += 1
        if self.nth == k:
            self.result = root.val
        self.dfs(root.right, k)

    """

    """ template:
TreeNode pNode = root;
while (!s.isEmpty() || pNode != null) {
    if (pNode != null) {
        s.push(pNode);
        pNode = pNode.left;
    } else {
        pNode = s.pop();
        result.add(pNode.val);
        pNode = pNode.right;
    }
}
    """
    def kthSmallest(self, root, k):
        if not root:
            return None
        q = []
        node = root
        nth = 0
        while q or node:
            if node:
                q.append(node)
                node = node.left
            else:
                node = q.pop()
                nth += 1
                if nth == k:
                    return node.val
                node = node.right
        return None

　　

86. 二叉查找树迭代器

中文

English

设计实现一个带有下列属性的二叉查找树的迭代器：
next()返回BST中下一个最小的元素

• 元素按照递增的顺序被访问（比如中序遍历）
• next()和hasNext()的询问操作要求均摊时间复杂度是O(1)

样例

样例 1:

输入：{10,1,11,#,6,#,12}
输出：[1, 6, 10, 11, 12]
解释：
二叉查找树如下 :
  10
  /\
 1 11
  \  \
   6  12
可以返回二叉查找树的中序遍历 [1, 6, 10, 11, 12]

样例 2:

输入：{2,1,3}
输出：[1,2,3]
解释：
二叉查找树如下 :
  2
 / \
1   3
可以返回二叉查找树的中序遍历 [1,2,3]

挑战

额外空间复杂度是O(h)，其中h是这棵树的高度

Super Star：使用O(1)的额外空间复杂度

"""
Definition of TreeNode:
class TreeNode:
    def __init__(self, val):
        self.val = val
        self.left, self.right = None, None

Example of iterate a tree:
iterator = BSTIterator(root)
while iterator.hasNext():
    node = iterator.next()
    do something for node
"""

class BSTIterator:
    """
    @param: root: The root of binary tree.
    """
    def __init__(self, root):
        # do intialization if necessary
        self.q = []
        self.node = root

    """
    @return: True if there has next node, or false
    """
    def hasNext(self, ):
        # write your code here
        return self.node or self.q

    """
    @return: return next node
    """
    def next(self, ):
        # write your code here
        while self.node or self.q:
            if self.node:
                self.q.append(self.node)
                self.node = self.node.left
            else:
                cur_node = self.q.pop()
                self.node = cur_node.right
                return cur_node
        return None