job 인터뷰/코테(Matroid) 준비

173. Binary Search Tree Iterator

hyunkookim 2024. 12. 14. 17:51

173. Binary Search Tree Iterator

📝 문제 설명

Implement the BSTIterator class that represents an iterator over the in-order traversal of a binary search tree (BST):
BSTIterator 클래스를 구현하세요. 
이 클래스는 **이진 탐색 트리(BST)의 중위 순회(in-order traversal)**를 따르는 이터레이터(iterator) 역할을 합니다.

BSTIterator(TreeNode root) Initializes an object of the BSTIterator class.
**BSTIterator(TreeNode root)**는 BSTIterator 클래스의 객체를 초기화합니다.

The root of the BST is given as part of the constructor.
생성자에서 BST의 루트 노드가 주어집니다.

The pointer should be initialized to a non-existent number smaller than any element in the BST.
포인터는 BST의 어떤 값보다도 작은 값을 가리키도록 존재하지 않는 가상의 값으로 초기화해야 합니다.
즉, 아직 아무 노드도 방문하지 않은 상태로 시작한다는 뜻입니다.

boolean hasNext() Returns true if there exists a number in the traversal to the right of the pointer, otherwise returns false.
hasNext() 함수는 현재 포인터 오른쪽에 순회할 값이 존재하면 true를 반환하고, 더 이상 값이 없으면 false를 반환합니다.

int next() Moves the pointer to the right, then returns the number at the pointer.
next() 함수는 포인터를 오른쪽으로 이동시킨 후, 이동한 위치의 값을 반환합니다.

Notice that by initializing the pointer to a non-existent smallest number, the first call to next() will return the smallest element in the BST.
포인터가 "존재하지 않는 가장 작은 값"에서 시작했기 때문에,
첫 번째 next() 호출은 BST에서 가장 작은 값을 반환하게 됩니다.

 

📌 Constraints (제약 조건)

The number of nodes in the tree is in the range [1, 10⁵].
→ 트리의 노드 개수는 최소 1개 이상, 최대 10만 개입니다.

0 <= Node.val <= 10⁶
→ 각 노드의 값은 0 이상 100만 이하의 정수입니다.

 

At most 10⁵ calls will be made to hasNext, and next.
→ hasNext()와 next() 함수는 최대 10만 번까지 호출될 수 있습니다.


🧠 Follow-up (추가 도전 조건)

Could you implement next() and hasNext() to run in average O(1) time and use O(h) memory, where h is the height of the tree?

다음 조건을 만족하는 방식으로 next()와 hasNext()를 구현할 수 있나요?

🎯 조건 해석

  • ✅ next()와 hasNext()는 평균 시간복잡도 O(1) 으로 동작해야 함
  • ✅ 메모리는 O(h) 만 사용해야 함 → h는 트리의 높이

🔍 무슨 의미인가요?

✅ 왜 O(1) 시간이어야 하나요?

  • next()는 중위 순회의 다음 값을 빠르게 반환해야 하므로,
    전체 배열을 매번 만들면 O(n)이 됩니다 ❌
  • 대신, 스택에 경로만 저장하면서 다음 노드만 찾도록 하면
    평균적으로 O(1) 에 처리할 수 있어요 (각 노드는 한 번만 스택에 push/pop 되니까요)

✅ 왜 O(h) 메모리인가요?

  • 전체 노드를 저장하면 O(n) 이므로 ❌
  • 현재 위치에서 다음으로 가기 위한 경로만 저장하면, 최대 깊이 h만큼만 저장하면 되므로 O(h)

🧩 핵심 아이디어 요약

  • inorder 순회를 미리 다 하지 말고,
  • 스택을 사용해 중위 순회를 "하나씩" 진행
  • 현재 방문한 노드 기준으로 다음 왼쪽 or 오른쪽 서브트리를 즉석에서 탐색

✅ 이런 구조로 구현해야 합니다

class BSTIterator:
    def __init__(self, root: Optional[TreeNode]):
        self.stack = []
        self._leftmost_inorder(root)

    def _leftmost_inorder(self, node):
        # 현재 노드에서 가장 왼쪽까지 스택에 쌓음
        while node:
            self.stack.append(node)
            node = node.left

    def next(self) -> int:
        # 스택에서 하나 꺼냄 (현재 가장 작은 값)
        top_node = self.stack.pop()

        # 오른쪽 서브트리가 있으면 다시 왼쪽 끝까지 스택에 쌓음
        if top_node.right:
            self._leftmost_inorder(top_node.right)

        return top_node.val

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

✅ 이 구현은

  • next()와 hasNext() 모두 평균 O(1) 시간
    (각 노드는 한 번만 push/pop 되니까)
  • 스택의 최대 크기 = 트리의 높이 h → O(h) 메모리 사용
# 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 = []
        self.saveMostLeft(root)
        
    def saveMostLeft(self, node): 
        # 구현한 코드!! 중요
        while node:
            self.stack.append(node)
            node = node.left    

    def next(self) -> int:
        # self.stack 의 가장 위에 있는 값이, 가장 왼쪽값임
        node = self.stack.pop()

        if node.right: # 오른쪽 노드가 있으면
            self.saveMostLeft(node.right) # 오늘쪽 노드값 추가하고
        return node.val # 현재 노드값 반환하고.

    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()

 



https://youtu.be/RXy5RzGF5wo?si=EcNVv54-gR9Tln1F

 

이 클래스는 BST를 중위 순회(in-order traversal) 순서로 순회하며, 노드 값을 반환하는 이터레이터를 구현합니다.
* in-order traversal: 중위 순회는 왼쪽 → 현재 노드 → 오른쪽 순서로 노드를 방문합니다.
클래스는 다음 두 가지 주요 메서드를 제공합니다:
next(): 이터레이터의 포인터를 다음 노드로 이동하고, 해당 노드의 값을 반환.
hasNext(): 이터레이터가 더 탐색할 노드가 있는지 확인.
next()와 hasNext()의 동작

 

DFS, 재귀로..
그러나, 제한 사항이..

  • next()와 hasNext()를 평균 O(1) 시간 안에 실행하고 
  • O(h) 메모리를 사용하도록 구현할 수 있나요? 
  • 여기서 h는 트리의 높이입니다.

따라서, 스텍으로 구현

 

# 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:
        """
        next()**는 현재 포인터를 이동시켜 다음 노드의 값을 반환
            중위 순회를 따라 왼쪽에서 오른쪽으로 이동
        스텍으로 구현하므로, 다음노드는 이미 스텍에 쌓여있음.
            현재 인덱스 바로 밑의 값임
        """
        res = self.stack.pop()

        # 이때 반환 전에, 오른쪽에 노드 있으면, 
        # 오른쪽 노드의 왼쪽 밑으로 쭉.. 내려가서.A
        # 미리 추가 하고..
        cur = res.right
        while cur:
            self.stack.append(cur)
            cur = cur.left

        return res.val

    def hasNext(self) -> bool:
        """
        hasNext()는 현재 포인터 기준으로 
            더 탐색할 노드가 있으면 True, 없으면 False를 반환
        stack 으로 구현하므로, hashNext call 했을때,
            stack 이 비워있으면, 다음 탐색 노드가 없다는 의미니, False
            아니면, True
        """
        return self.stack != []  # 비어있지 않으면 True
        


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