BST 方法，返回指定范围内的值列表 Python 实现

Question

我想返回一个排序顺序列表，前提是我获得了该方法的开始/停止值。例如，如果 start=2 和 end=8，那么我想隐式返回该范围内的 BST 中按排序顺序的值的列表。

由于我希望它按排序顺序并且不允许在方法调用之后对列表进行排序，所以我认为我应该通过顺序遍历来遍历 bst。当我测试我的实现时，首先第一个 doctest 按预期返回 [7,9,11] 而不是 [5,7,9,11]。

from __future__ import annotations
from typing import Any, List, Optional, Tuple

class BinarySearchTree:
    """Binary Search Tree class.

    # === Representation Invariants ===
    #  - If self._root is None, then so are self._left and self._right.
    #    This represents an empty BST.
    #  - If self._root is not None, then self._left and self._right
    #    are BinarySearchTrees.
    #  - (BST Property) If self is not empty, then
    #    all items in self._left are <= self._root, and
    #    all items in self._right are >= self._root.
    """
    def __init__(self, root: Optional[Any]) -> None:
        """Initialize a new BST containing only the given root value.

        If <root> is None, initialize an empty tree.
        """
        if root is None:
            self._root = None
            self._left = None
            self._right = None
        else:
            self._root = root
            self._left = BinarySearchTree(None)
            self._right = BinarySearchTree(None)

    def is_empty(self) -> bool:
        """Return True if this BST is empty.

        >>> bst = BinarySearchTree(None)
        >>> bst.is_empty()
        True
        >>> bst = BinarySearchTree(10)
        >>> bst.is_empty()
        False
        """
        return self._root is None

    def items_in_range(self, start: Any, end: Any) -> List:
        """Return the items in this BST between <start> and <end>, inclusive.

        Precondition: all items in this BST can be compared with <start> and
        <end>.
        The items should be returned in sorted order.

        As usual, use the BST property to minimize the number of recursive
        calls.

        >>> bst = BinarySearchTree(7)
        >>> left = BinarySearchTree(3)
        >>> left._left = BinarySearchTree(2)
        >>> left._right = BinarySearchTree(5)
        >>> right = BinarySearchTree(11)
        >>> right._left = BinarySearchTree(9)
        >>> right._right = BinarySearchTree(13)
        >>> bst._left = left
        >>> bst._right = right
        >>> bst.items_in_range(4, 11)
        [5, 7, 9, 11]
        >>> bst.items_in_range(10, 13)
        [11, 13]
        """
        if self.is_empty():
            return []
        else:
            #use helper here
            if end >= self._root >= start:
                return (self._left._helper_items_in_range_left(start)
                        + [self._root]
                        + self._right._helper_item_in_range_right(end))
            elif self._root > end:
                return self._left.items_in_range(start,end)
            elif self._root < start:
                return self._right.items_in_range(start,end)
            else:
                pass

    def _helper_items_in_range_left(self, start):
        if self.is_empty():
            return []
        elif self._root < start:
            return []
        else:
            return self._left._helper_items_in_range_left(start) +\
                   [self._root] + self._right._helper_items_in_range_left(start)

    def _helper_item_in_range_right(self, end):
        if self.is_empty():
            return []
        elif self._root > end:
            return []
        else:
            return self._left._helper_item_in_range_right(end) + [self._root] +\
                   self._right._helper_item_in_range_right(end)

Answer 1

你可以使用这样的东西。请注意，我使用不同的树结构对其进行了测试。

import itertools
from collections import deque
class BSTIterator(object):

    def __init__(self, root):
        # Constructor takes in a tree root
        self.stack = deque()
        self._get_min(root)

    def _get_min(self, root):
        # We need to create our stack, i.e. dive down the left
        curr = root
        while curr != None:
            self.stack.append(curr)
            curr = curr.left        

    def __iter__(self): # Return self as the iterator
        return self
    def __next__(self): # Every time `next` is called return our data.
        try:
            curr = self.stack.pop()
            self._get_min(curr.right)
            return curr.data
        except IndexError:
            raise StopIteration

使用的树类型：

class Node:
    def __init__(self, data):
        self.left = None
        self.right = None
        self.data = data
    def insert(self, data):
        if self.data:
            if data < self.data:
                if self.left is None:
                    self.left = Node(data)
                else:
                    self.left.insert(data)
            elif data > self.data:
                if self.right is None:
                    self.right = Node(data)
                else:
                    self.right.insert(data)
        else:
            self.data = data

经测试：

root = Node(8)
root.insert(3)
root.insert(10)
root.insert(1)
root.insert(7)
root.insert(12)
root.insert(121)
root.insert(23)
root.insert(19)
root.insert(9)
b_iter = BSTIterator(root)
# root.print_tree()

# Since we now have an iterator we can for loop over it
# such as
# y = [x for x in b_iter]
# or we can slice it like
y = [x for x in itertools.islice(b_iter, 2, 5)]
print(y)

打印：

[7, 8, 9]