[알고리즘/python] LeetCode 230. Kth Smallest Element in a BST - iterative/recursive 풀이, follow up 답변 포함

📑 1. 문제 이해

https://leetcode.com/problems/kth-smallest-element-in-a-bst/description/?envType=study-plan-v2&envId=top-interview-150 

Kth Smallest Element in a BST - LeetCode

이 문제는 이진 탐색 트리에서 k번째 작은 수를 찾는 것입니다. 이때 주의할 점은 다음과 같습니다. 

- 빈 트리는 주어지지 않습니다. root는 항상 None이 아닙니다.

- 노드의 값은 항상 0보다 크거나 같습니다.

- k가 n보다 크거나, 1보다 작은 비정상적인 경우는 입력으로 주어지지 않습니다.

☑️ 테스트케이스

문제에서 주어진 테스트케이스는 다음과 같습니다.

 

Example1: root = [3, 1, 4, null, 2], k = 1

- 찾는 노드가 최하위 레벨에 존재합니다. 부모 노드의 오른쪽에 위치한 노드입니다.

 

Example2: root = [5, 3, 6, 2, 4, null, null, 1], k = 3

- 찾는 노드가 내부 노드에 위치합니다. 

 

다음과 같은 경우를 추가로 살펴볼 필요가 있습니다.

 

root = [1], k =1

- 루트 노드만 존재합니다.

 

Input으로 주어진 배열이 실제로 함수의 매개변수로 어떻게 들어오는지 알고 싶다면  이전 발행 글 의 배열로 트리 만들기를 참고해주세요.

💡 2. 알고리즘 설계

이진 탐색 트리의 중위 순회하면 오름차순으로 정렬된 결과를 얻을 수 있습니다. 반복문을 통한 깊이 우선 탐색 - 중위 순회알고리즘을 사용하여 k번째 작은 원소를 찾을때까지 탐색합니다. 

    def kthSmallest(self, root: TreeNode, k: int):
        stack = []
        node = root
        count = 0
        while stack or node:
            while node:
                stack.append(node)
                node = node.left

            count += 1
            node = stack.pop()
            if count == k:
                return node.val

            node = node.right

시간 복잡도 및 공간 복잡도

시간 복잡도: O(H + k). H는 BST의 높이이다. 왼쪽 자식을 모두 스택에 추가할때까지 while문이 실행되므로 LogN번(트리의 높이)만큼 실행된다. 이후 k번째 노드를 찾을때까지 거슬러 올라오면서 탐색한다.

 

공간 복잡도: O(LogN).  스택에는 최대 H(트리의 높이)개의 노드가 담기는데, 트리의 높이 H는 최악의 경우 N, 균형잡힌 트리의 경우 LogN이 된다. 연결리스트 형태가 되는 경우는 매우 드물기 때문에 분할 상환 공간 복잡도(Amortized)는 O(LogN)이 된다.

 

⚒️ 3. 최적화

이 문제에는 follow up 질문이 있습니다.

Follow up:
 If the BST is modified often (i.e., we can do insert and delete operations) and you need to find the kth smallest frequently, how would you optimize?

"만약 BST가 빈번하게 삽입 삭제되는 상황에서 k번째 원소를 자주 찾아야한다면 어떻게 최적화할 수 있을것인가"에 대한 질문입니다. 

 

 DB 인덱스에서처럼 자주 삽입 삭제되는 곳에서 사용되는 트리를 생각해볼 수 있습니다. 이때 중요한 것은 트리의 높이를 LogN으로 제한하는 균형 트리(AVL, Red-Black Tree등)를 사용하는 것입니다.

 

 또한,  트리에 왼쪽 서브 트리에 속하는 노드의 개수를 기록하면 k번째 원소를 찾는 탐색과정을 최적화할 수 있습니다. 왼쪽 서브트리의 노드의 개수보다 k가 작다면 루트 노드의 오른쪽 서브트리에서 찾습니다. 이외에는 왼쪽 서브 트리에서 찾습니다. k를 줄여가며 서브트리를 탐색해갑니다. 서브 트리 중 한쪽만 선택해서 탐색할 수 있게 되어 불필요한 연산이 줄어듭니다. 이때 주의할 점은 삽입, 삭제 시에 부모 노드들에 기록된 왼쪽 서브트리 노드의 개수를 조정해주어야한다는 것입니다. 

 

코드는 다음과 같습니다. 

def kth_smallest(root, k):
    if not root:
        return None
    
    # Calculate the rank of the current node
    rank_of_root = 1 + (root.left.rank if root.left else 0)
    
    # Case 1: If k is equal to the rank of the current node, we've found the kth smallest element
    if k == rank_of_root:
        return root.val
    
    # Case 2: If k is less than the rank of the current node, search in the left subtree
    elif k < rank_of_root:
        return kth_smallest(root.left, k)
    
    # Case 3: If k is greater than the rank of the current node, search in the right subtree
    else:
        return kth_smallest(root.right, k - rank_of_root)
        
def insert(root, val):
    if not root:
        return TreeNode(val)
    
    if val < root.val:
        root.left = insert(root.left, val)
        root.rank += 1  # Increment rank for the left child
    else:
        root.right = insert(root.right, val)
    
    return root
    
  
 def delete(root, key):
    if not root:
        return root
    
    # Perform the regular deletion operation
    if key < root.val:
        root.left = delete(root.left, key)
    elif key > root.val:
        root.right = delete(root.right, key)
    else:
        # Node with the key to be deleted found
        if not root.left:
            return root.right
        elif not root.right:
            return root.left
        
        # Node with two children: Get the in-order successor
        successor = find_min(root.right)
        root.val = successor.val
        root.right = delete(root.right, successor.val)
    
    # Update ranks during traversal
    if root.left:
        root.rank = 1 + root.left.rank
    else:
        root.rank = 0
    
    return root

---

🤔 다른 사람의 풀이

재귀 함수를 이용하여 DFS를 수행할 수도 있습니다.  다만, 위의 반복문을 활용한 DFS와 같이 k개의 노드에 대해 중위 순회를 해야함으로 실행 시간은 차이나지 않습니다. 

   이때 주의할 점은 이미 k번째 수를 찾은 경우에는 오른쪽 자식에 대해 DFS를 수행하지 않도록 해 불필요한 연산을 줄여야합니다.

    def kthSmallest_recursive(self, root: TreeNode, k: int):
        count = 0
        result = None

        def __kthSmallest(node: TreeNode):
            nonlocal count, k, result

            if node.left:
                __kthSmallest(node.left)

            count += 1
            if count == k:
                result = node.val

            if count < k and node.right:
                __kthSmallest(node.right)

        __kthSmallest(root)
        return result

---

✅ 정리

◎ 이진 탐색 트리를 효율적으로 사용하기 위해서는 균형 트리를 사용해야합니다. 

◎ 깊이 우선 탐색(DFS)의 경우 반복문과 스택을 활용한 버전과 재귀 버전이 있습니다.