스위프트로 구현하는 자료구조 8: 이진 탐색 트리(Binary Search Tree)

이진 탐색 트리 (BST)

이진 탐색 트리는 이진 트리에 몇가지 조건이 더 추가된 트리입니다. 이진 탐색 트리는 다음 조건을 만족해야합니다.

 

1) 어떤 루트노드(서브트리의 루트 포함)의 왼쪽 서브트리의 노드들이 가진 값은 루트 노드가 가진 값보다 작아야합니다.

2) 어떤 루트노드의 오른쪽 서브트리의 노드들이 가진 값은 루트 노드가 가진 값보다 커야합니다.

3) 노드의 값으로 트리 연산이 이루어지기 때문에 중복된 노드 값을 허용하지 않습니다. 

예를 들어, 이런 트리는 이진 탐색 트리가 될 수 없습니다. 루트 노드인 1보다 큰 값인 2가 왼쪽 자식 노드에 위치해있고 왼쪽 자식 노드인 4, 6도 모두 자신의 부모 노드의 값보다 크기 때문입니다

반면에 이런 트리는 이진 탐색 트리가 될 수 있습니다. 루트노드인 4의 왼쪽 서브트리의 모든 노드들은 4보다 작고, 오른쪽 서브트리의 모든 노드들은 4보다 큽니다. 그리고 노드 2와 6 역시도 왼쪽에는 자신보다 작은 값, 오른쪽에는 자신보다 큰 값을 가지고 있습니다.

 

뼈대 만들기

TreeNode

import Foundation

final class TreeNode<T: Comparable> {
    private(set) var data: T
    var leftChild: TreeNode?
    var rightChild: TreeNode?
    
    var isLeafNode: Bool {
        return self.leftChild == nil && self.rightChild == nil
    }
    var hasSingleChild: Bool {
        return (self.leftChild != nil && self.rightChild == nil)
        || (self.leftChild == nil && self.rightChild != nil)
    }
    var hasTwoChild: Bool {
        return self.leftChild != nil && self.rightChild != nil
    }
    
    init(data: T) {
        self.data = data
    }
    
    func update(data: T) {
        self.data = data
    }
    
    var asString:String { return treeString(self){("\($0.data)",$0.leftChild,$0.rightChild)}  }
}

먼저 트리에 사용될 노드를 만들어주었습니다. 가장 중요한 것은 데이터를 저장하는 변수, 왼쪽 자식 노드에 대한 포인터, 그리고 오른쪽 자식 노드에 대한 포인터입니다.

 

그리고 노드의 상태를 체크할 수 있는 세 가지 프로퍼티를 만들었습니다. 노드의 자식노드 유무에 따라서 자식 노드가 없는 리프 노드인지, 자식 노드가 하나인지, 두 개인지 확인해주게 됩니다. 이런 프로퍼티가 왜 필요한지는 잠시후에 다시 설명하도록 하겠습니다. 

 

추가적으로 현재 데이터를 새로운 데이터로 업데이트하는 메서드도 만들어주었는데요, 이 메서드 역시 어떻게 사용되는지 잠시후에 설명할게요!

 

asString 프로퍼티는 트리를 엄청 예쁘게 출력해주는 treeString 메서드를 호출하는데 이건 스택 오버플로우에서 진짜 잘 만들어주신 분이 계셔서 그대로 사용했어요!

 

이 질문에 대한 답변에 메서드가 정의되어 있구요,

How to “draw” a Binary Tree in Console?

 

적용하면 이렇게 예쁘게 콘솔에 트리가 출력됩니다!

탐색(Search) - O(log N) 

이진 탐색 트리의 특징이 어떤 노드의 왼쪽 서브트리는 항상 자신보다 작은 값들이 있고, 오른쪽에는 항상 자신보다 큰 값들이 있다는 것이었습니다. 탐색도 이를 이용해서 진행합니다.

이 트리에서 5를 값으로 가진 노드를 찾아가려면 어떻게 해야할까요? 트리이니까 일단 루트노드에서 시작해야겠죠?

 

루트 노드의 값인 4와 찾고자 하는 노드의 값인 5를 비교해봅시다. 찾고자 하는 노드의 더 값이 크네요. 그럼 이진 탐색 트리의 규칙대로 왼쪽 서브 트리는 탐색할 필요없이 오른쪽 서브트리로 이동합시다!

 

이번엔 새로운 루트 노드와 그 값인 6을 만나게됩니다. 6은 찾고자 하는 노드의 값보다 크기 때문에 현재 루트노드를 기준으로 오른쪽 서브트리는 탐색할 필요가 없습니다. 왼쪽 서브트리로 이동할게요!

이번에는 새로운 루트 노드의 값이 찾고자하는 값과 같네요! 그럼 이 노드의 참조값을 그대로 반환해주면 원하는 노드의 탐색을 마치게 됩니다. 만약 이번 노드의 값이 5가 아니었다면 nil을 만나면서 탐색이 종료되겠죠?

 

위 과정을 코드로 적으면 이렇습니다.

    private func search(data: T, in node: TreeNode<T>?) -> TreeNode<T>? {
        if node == nil || node?.data == data { return node }
        
        return node!.data > data
        ? search(data: data, in: node?.leftChild)
        : search(data: data, in: node?.rightChild)
    }

현재 루트 노드(node)의 데이터가 찾는 데이터의 값보다 크면 왼쪽 서브트리로 재귀하고, 작으면 오른쪽 서브트리로 재귀합니다. 그리고 현재 루트노드가 nil이거나, 두 데이터가 일치하면 재귀를 종료하고 현재 루트노드를 그대로 반환합니다. 

 

탐색하는 과정에서 트리의 한쪽 서브트리의 탐색을 계속 생략하면서 진행하기 때문에 시간복잡도는 트리의 높이에 따른 노드개수인 O(log N)으로 정리됩니다.

 

탐색  Worst Case - O(N)

하지만 만약 이진 탐색 트리가 다음처럼 구성되어 있다면 어떨까요?

각 서브트리의 루트노드의 왼쪽에는 자신보다 작은 값들만 있기 때문에 이진 탐색 트리의 조건에는 위배되지 않지만 편향 이진 트리가 되어버립니다. 이런 상황에서 탐색을 진행하면 계속해서 왼쪽 서브트리로만 이동하게 되고 생략하는 서브트리가 없기 때문에 결국 모든 노드를 다 확인해는 O(N)의 시간복잡도가 발생합니다. 

 

이런 문제를 막기 위해 BBST(Balanced Binary Search Tree)인 AVL 트리와 Red-Black 트리가 고안되었습니다. 나중에 기회가 된다면 구현해보겠습니다! 결론은 일반 이진 탐색 트리로는 이런 문제를 막을 수 없습니다..

 

삽입(add) - O(log N)

거의 모든 이진 트리의 연산들은 탐색에 의존합니다. 삽입 역시도 적절한 위치에 새로운 값을 가진 노드를 넣어주기 위해서 탐색을 진행해야합니다.

이번에는 이렇게 생긴 트리에 새로운 노드를 추가한다고 해보겠습니다.

이번에 추가할 노드의 값은 4입니다. 루트 노드와 먼저 값을 비교해봅시다. 새 노드의 값이 더 작네요. 왼쪽 서브트리로 들어갑시다.

이번엔 새로운 노드의 값이 서브트리의 루트노드 보다 더 큰 것을 확인했습니다. 그럼 오른쪽 서브트리로 가야겠죠?

새로운 노드의 값이 서브트리의 루트노드보다 큽니다. 이제 오른쪽으로 이동하려고 하는데, 현재 루트 노드는 리프 노드라서 오른쪽에 자식 노드가 존재하지 않습니다. 따라서 새로운 노드를 오른쪽 자식 노드로 만들어주면 새로운 노드의 추가가 완료됩니다.

이 과정에서 본 것처럼 실제로 새로운 노드를 트리에 추가하는 것은 상수시간이지만 새로운 위치를 탐색하는데까지 탐색하는 시간이 요구되기 때문에 시간복잡도는 탐색의 시간복잡도와 동일한 O(log N)이 됩니다.

 

    mutating func add(data: T) {
        let node = TreeNode(data: data)
        if root == nil {
            self.root = node
            return
        }
        self.add(newNode: node, under: root)
    }

    @discardableResult
    private func add(newNode: TreeNode<T>, under parentNode: TreeNode<T>?) -> TreeNode<T>? {
        guard newNode.data != parentNode?.data else { return nil }
        if parentNode == nil {
            return newNode
        }
        parentNode!.data > newNode.data
        ? (parentNode?.leftChild = add(newNode: newNode, under: parentNode?.leftChild))
        : (parentNode?.rightChild = add(newNode: newNode, under: parentNode?.rightChild))
        return parentNode
    }

구현은 탐색과 동일하게 진행되지만 리프 노드의 자식 노드를 만나면 새로운 노드를 삽입해주는 로직이 추가됩니다. 

삭제 - O(log N)

BST의 최고 헷갈리는 포인트는 삭제 연산이죠ㅠ 삭제할 노드의 자식 노드 개수에 따라서 총 세 개의 경우로 로직이 다르게 구현됩니다. 

 

삭제 연산도 삭제할 노드를 먼저 찾아내야 하기 때문에 탐색이 선행됩니다. 탐색이 완료되어서 삭제할 노드를 찾으면 다음 과정을 진행합니다.

 

1) 삭제할 노드가 리프노드인 경우 

 

가장 쉬운 경우입니다. 리프 노드의 경우에는 왼쪽이나 오른쪽에 다른 노드들이 없기 때문에 그냥 트리에서 리프노드만 제거해주면 삭제가 완료됩니다.

간단하죠? 코드로 작성해보면 이렇게 됩니다.

        guard let targetNode = self.search(data: data, in: root) else { return nil }
	
    	// target is leaf node
        if targetNode.isLeafNode {
            let parent = self.findParent(of: targetNode, under: root)
            parent?.rightChild === targetNode
            ? (parent?.rightChild = nil)
            : (parent?.leftChild = nil)
            return targetNode.data
        }

삭제할 노드를 찾고 해당 노드의 부모 노드가 가진 삭제할 노드에 대한 포인터를 지워주면 됩니다. 사실 탐색하면서 부모노드를 찾는게 가장 좋은 방법인데 편의상 미리 구현된 탐색 알고리즘을 그대로 사용하고 싶어서 부모 노드를 찾는 메서드를 따로 만들어주었습니다.

 

2) 삭제할 노드에게 하나의 자식 노드가 있는 경우

이번에는 삭제할 노드가 자식 노드를 왼쪽이나 오른쪽에 하나만 가지고 있는 경우입니다. 이대로 노드를 삭제하면 아래에 딸려있던 서브 트리에 속하는 모든 노드가 삭제되기 때문에 삭제할 노드가 가진 하나의 서브트리를 기존 트리에 연결시켜주어야합니다.

왼쪽에 자식노드를 가지고 있는 값이 2인 노드를 삭제해봅시다.

삭제할 노드의 부모 노드를 찾고 해당 노드의 자식 노드를 삭제할 노드의 자식 노드로 변경해줍니다. 이진 탐색 트리이기 때문에 이미 왼쪽 서브트리는 모두 삭제할 노드의 부모 노드보다 작은 값들을 가지고 있기 때문에 따로 트리를 재조정해 줄 필요가 없습니다.

 

        guard let targetNode = self.search(data: data, in: root) else { return nil }
       
       // target has single child node 
        if targetNode.hasSingleChild {
            let child = targetNode.leftChild != nil ? targetNode.leftChild : targetNode.rightChild
            guard let parent = self.findParent(of: targetNode, under: root) else { return nil }
            parent.data > child!.data ? (parent.leftChild = child) : (parent.rightChild = child)
            return targetNode.data
        }

코드로 구현할 때도 어렵지 않게 구현할 수 있었습니다. 삭제할 노드의 부모 노드를 찾아 자식 노드 정보를 변경해주는 것으로 쉽게 구현이 가능합니다.

 

3) 삭제할 노드에게 두 개의 자식 노드가 있는 경우

문제는 삭제할 노드가 두 개의 자식노드를 가지고 있는 경우입니다.

이번에는 이 트리의 루트 노드를 삭제한다고 해보겠습니다. 루트노드를 삭제하면서 양쪽 서브트리들을 연결해야되는데요, 어떻게 해야 이진 탐색 트리를 유지하면서 연결해줄 수 있을까요?

 

삭제될 노드의 양쪽 서브트리를 연결할 노드는 왼쪽 서브트리의 값보다는 커야하고, 오른쪽 서브트리의 값보다는 작아야합니다. 

 

이 조건을 만족하는 노드는 오른쪽 서브트리 가장 왼쪽에 있는 노드를 찾아 얻을 수 있습니다. 오른쪽 서브트리는 항상 왼쪽 서브트리보다 값이 크고, 오른쪽 서브트리의 가장 왼쪽 서브트리는 그 중에서도 가장 작은 값이기 때문에 왼쪽 서브트리의 리프 노드를 사용하면 이진 탐색 트리를 유지하면서 양쪽 서브트리를 연결시켜줄 수 있습니다.

 

이렇게 삭제할 노드의 오른쪽 자식의 서브트리에 있는 가장 작은 값을 가진 노드를 Successor 라고 합니다. 삭제할 노드의 위치를 계승받을 노드라는 의미인 것 같네요.

 

삭제할 노드의 오른쪽 자식 노드로 먼저 이동하고, 해당 노드의 왼쪽 서브트리를 계속 따라가면,

값이 7인 노드를 얻을 수 있습니다. 

 

이제 이 노드가 가진 값을 루트노드에 덮어씌워줍시다.

실제로 노드를 교환하는 것도 좋겠지만 키로 사용되는 노드의 값으로 덮어씌워도 무방합니다. 만약 노드 인스턴스를 삭제 이후에 반환하게 한다면 노드 자체를 교환해야겠죠?

그리고 Successor로 사용된 노드를 지워주면, 목표 노드를 제거하고도 이진 탐색 트리를 유지할 수 있습니다.

 

        // target has two child nodes
        if targetNode.hasTwoChild {
            let deleted = targetNode.data
            guard let successor = self.findSuccessor(of: targetNode.rightChild) else { return nil }
            if !successor.isLeafNode {
                let parent = self.findParent(of: successor, under: root)
                parent?.leftChild = successor.rightChild
            }
            _ = self.delete(data: successor.data)
            targetNode.update(data: successor.data)
            return deleted
        }

코드로는 이렇게 구현됩니다. 먼저 Successor를 찾아서 값을 덮어씌운 뒤에, Successor의 부모노드를 찾아서 자식 노드의 참조를 제거해줍니다.

 

이렇게 해서 삭제를 정리했는데요, 삭제도 탐색과 삭제의 조합이기 때문에 O(logN)의 시간복잡도가 요구됩니다. 제 구현에서는 부모 노드를 찾을 때 탐색을 한번 더 진행해서 개인적으로는 좋지 않다고 생각하는데, 부모 노드에 대한 정보를 노드에 프로퍼티로 가지고 있거나, 탐색할 때 부모노드의 정보까지 반환해주면 최적화가 가능할 것 같습니다.

 

전체 코드

 

GitHub - jeonyeohun/Data-Structures-In-Swift: 👷🏻‍♂️ STL 왜 없어.. 스위프트로 구현하는 자료구조

자료구조 시리즈의 코드들은 이 레포에서 모두 확인할 수 있어요!