Skip to content

173. Binary Search Tree Iterator

#211

Problem link

https://leetcode.com/problems/binary-search-tree-iterator/

Problem Summary

BST에서 크기 순으로 동작하는 이터레이터를 만드는 문제. 메소드는 next(), hasNext() 두개를 구현하면 된다.

Solution

초기화 때 inorder으로 전부 돌아서 배열에 넣은 다음 하나씩 출력하는 방법이 간단하지만 메모리는 O(n)이긴 하다.
스택을 쓰면 O(h)로 메모리를 줄일 수 있다.

하지만 결과에서의 실행 시간은 배열에 전부 집어넣은 방식이 더 빠르고 메모리도 적게 사용한다... 오차가 있을 수 있겠지만 단순 2번씩 제출 결과는 그렇다.

Source Code

Inorder Array

  • init time: O(n)
  • next() time: O(1)
  • space: O(n)
  • result: 71ms
from typing import Optional


# 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 BSTIterator:

    def __init__(self, root: Optional[TreeNode]):
        self.arr = []
        self.idx = 0

        def inorder(node: Optional[TreeNode]):
            if node is None:
                return
            inorder(node.left)
            self.arr.append(node.val)
            inorder(node.right)

        inorder(root)

    def next(self) -> int:
        ret = self.arr[self.idx]
        self.idx += 1
        return ret

    def hasNext(self) -> bool:
        return self.idx < len(self.arr)

# Your BSTIterator object will be instantiated and called as such:
# obj = BSTIterator(root)
# param_1 = obj.next()
# param_2 = obj.hasNext()

Inorder Stack

  • init time: O(h)
  • next() time: O(1) (average)
  • space: O(h)
  • result: 125ms
from typing import Optional


# 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 BSTIterator:

    def __init__(self, root: Optional[TreeNode]):
        self.stack = []
        while root:
            self.stack.append(root)
            root = root.left

    def next(self) -> int:
        node = self.stack.pop()
        ret = node.val

        if node.right:
            node = node.right
            while node:
                self.stack.append(node)
                node = node.left
        return ret

    def hasNext(self) -> bool:
        return len(self.stack) > 0

# Your BSTIterator object will be instantiated and called as such:
# obj = BSTIterator(root)
# param_1 = obj.next()
# param_2 = obj.hasNext()