LeetCode/주제별 보충

BST Sets and Maps ★★

hyunkookim 2025. 1. 16. 14:51

 

Design Binary Search Tree

Solved 

Medium

Design a Binary Search Tree class.

You will design a Tree Map, which maps an integer key to an integer value. Your Tree class should support the following operations:

  • TreeMap() will initialize an binary search tree map.
  • void insert(int key, int val) will map the key to the value and insert it into the tree.
  • int get(int key) will return the value mapped with the key. If the key is not present in the tree, return -1.
  • int getMin() will return the value mapped to the smallest key in the tree. If the tree is empty, return -1.
  • int getMax() will return the value mapped to the largest key in the tree. If the tree is empty, return -1.
  • void remove(int key) will remove the key-value pair with the given key from the tree.
  • int[] getInorderKeys() will return an array of the keys in the tree in ascending order.

Note:

  • The tree should be ordered by the keys.
  • The tree should not contain duplicate keys. If the key is already present in the tree, the original key-value pair should be overridden with the new key-value pair.

Example 1:

Input:
["insert", 1, 2, "get", 1, "insert", 4, 0, "getMin", "getMax"]

Output:
[null, 2, null, 2, 0]
 

Example 2:

Input:
["insert", 1, 2, "insert", 4, 2, "insert", 3, 7, "insert", 2, 1, "getInorderKeys", "remove", 1, "getInorderKeys"]

Output:
[null, null, null, null, [1, 2, 3, 4], null, [2, 3, 4]]
 

 

class TreeNode:
    def __init__(self, key, val):
        self.key = key
        self.val = val
        self.left = None
        self.right = None

class TreeMap:
    
    def __init__(self):
        self.root = None


    def insert(self, key: int, val: int) -> None:
        newNode = TreeNode(key, val) # 노드를 하나 새로 만들고

        if not self.root: # self.root 가 null 이면
            self.root = newNode # 만든 노드를 root 로 세팅하고
            return 

        curr = self.root
        while True:            
            if key < curr.key:
                if not curr.left:
                    curr.left = newNode
                    return
                curr = curr.left
            elif key > curr.key:
                if not curr.right:
                    curr.right = newNode
                    return
                curr = curr.right
            else: # key == curr.key
                curr.val = val
                return


    def get(self, key: int) -> int:        
        curr = self.root
        while curr:
            if key < curr.key:
                # if not curr.left:
                #     return -1
                curr = curr.left
            elif key > curr.key:
                # if not curr.right:
                #     return -1
                curr = curr.right
            else: # key == curr.key
                return curr.val 

        return -1

    def getMin(self) -> int:
        # curr = self.root
        # # if not curr:
        # #     return -1
        # while curr and curr.left:
        #     curr = curr.left
        curr = self.findMin(self.root)
        return curr.val if curr else -1

    def findMin(self, node) -> TreeNode:
        while node and node.left:
            node = node.left
        return node


    def getMax(self) -> int:
        curr = self.root
        # if not curr:
        #     return -1
        while curr and curr.right:
            curr = curr.right
        return curr.val if curr else -1


    def remove(self, key: int) -> None:
        # BST 에서 remove 구현은 재귀..가 가장 쉬운 솔루션.
        # 이를 위해서, helper 함수를 만들어서 진행하도록.
        self.root = self.removeHelper(self.root, key)

    # Remove the node with key, return the root of the subtree
    def removeHelper(self, curr, key) -> TreeNode:
        if curr == None:
            return None

        if key < curr.key:
            curr.left = self.removeHelper(curr.left, key)
        elif key > curr.key:
            curr.right = self.removeHelper(curr.right, key)
        else: # key == curr.key
            if not curr.right: # 오른쪽이 없으면, 왼쪽과 연결
                return curr.left
            elif not curr.left: # 왼쪽이 없으면, 오른쪽과 연결
                return curr.right
            else: #둘다 있으면
                # Swap curr node with inorder successor

                # 오른쪽의 가장 작은 키를 가진 노드를 찾아서
                minNode = self.findMin(curr.right)
                curr.key = minNode.key
                curr.val = minNode.val
                curr.right = self.removeHelper(curr.right, minNode.key)
        return curr


    def getInorderKeys(self) -> List[int]:
        result = []
        self.inorderTraversal(self.root, result)
        return result

    def inorderTraversal(self, root, result):
        # if not root:
        #     return result

        # inorder: 왼, 현, 오
        if root:
            if root.left:
                self.inorderTraversal(root.left, result)
            result.append(root.key)
            if root.right:
                self.inorderTraversal(root.right, result)