[Data Structure] 이진 탐색 트리(Binary Search Tree) 구현해보기 (feat.Swift)

안녕하세요 Foma 💻 입니다!

 

저번 시간에 이진 탐색 트리가 뭔지 이론에 대해 알아보았는데요.

 

오늘은 이진 탐색 트리를 직접 구현해보려고 합니다.

 

바로 시작할게요~

---

BSTNode

 

이진 탐색 트리의 노드를 만들어 주겠습니다.

 

노드는 값,왼쪽,오른쪽 자식의 노드 정보를 가지고 있습니다. 

 

class BSTNode<T:Comparable> {
    var value:T
    var leftChild:BSTNode?
    var rightChild:BSTNode?
    
    init(value:T) {
        self.value = value
    }
}

---

Binary Search Tree

 

이진 탐색 트리의 최상단 값을 가질 root를 프로퍼티로 정의했습니다.

 

struct BinarySearchTree<T:Comparable> {
    
    var root:BSTNode<T>?

---

Insert

 

public mutating func insert(value:T) {
        guard let _ = root else {
            self.root = BSTNode(value: value)
            return
        }
        
        var node:BSTNode = root!
        
        while true {
            if node.value > value {
                if node.leftChild == nil {
                    node.leftChild = BSTNode(value: value)
                    break
                }else {
                    node = node.leftChild!
                }
            }else {
                if node.rightChild == nil {
                    node.rightChild = BSTNode(value: value)
                    break
                }else {
                    node = node.rightChild!
                }
            }
        }
    }

 

값이 잘 들어가는지 테스트 해봐야겠죠?

 

테스트 하기 위해서 아래와 같이 왼쪽과 오른쪽 자식들을 재귀하며 출력하는 메서드를 구현했습니다.

 

 func printChild(node:BSTNode<Int>?) {
        if node == nil { return }
        print("\(node!.value)의 왼쪽 자식 :\(node!.leftChild?.value)")
        print("\(node!.value)의 오른쪽 자식 :\(node!.rightChild?.value)")
        printChild(node: node!.leftChild)
        printChild(node: node!.rightChild)
    }

 

10,11,8,1,3을 차례로 삽입하고 출력해보면 아래와 같이 값이 잘 들어간 것을 확인할 수 있습니다.

 

        var bst = BinarySearchTree<Int>()
        bst.insert(value: 10)
        bst.insert(value: 11)
        bst.insert(value: 8)
        bst.insert(value: 1)
        bst.insert(value: 3)
        printChild(node: bst.root)
        
//        10의 왼쪽 자식 :Optional(8)
//        10의 오른쪽 자식 :Optional(11)
//        8의 왼쪽 자식 :Optional(1)
//        8의 오른쪽 자식 :nil
//        1의 왼쪽 자식 :nil
//        1의 오른쪽 자식 :Optional(3)
//        3의 왼쪽 자식 :nil
//        3의 오른쪽 자식 :nil
//        11의 왼쪽 자식 :nil
//        11의 오른쪽 자식 :nil
        
        //        10
        //    8       11
        //1
        //    3

---

Search

 

'이진 탐색 트리의 검색은 어떤 걸 반환해야 할까?' 라고 고민했는데, 보통 해당 값이 있는지 없는지를 반환하더라구요.

 

그래서 저는 Search라는 이름 대신 isContain으로 이름지어줬습니다.

 

public func isContain(value:T) -> Bool {
        var node:BSTNode? = root
        while node != nil {
            if node!.value == value {
                return true
            }
            node = node!.value > value ? node?.leftChild : node?.rightChild
        }
        return false
    }

 

Insert에서 삽입한 값을 토대로 검색이 잘 되는지 확인해보겠습니다.

 

아래와 같이 포함되어 있다면 true를 그렇지 않다면 false를 반환하는 것을 볼 수 있습니다.

 

 bst.isContain(value: 11) //true
 bst.isContain(value: 100) //false
 bst.isContain(value: 0) //false
 bst.isContain(value: 8) //true

---

Delete

 

삭제를 구현할 때 진짜 복잡했습니다..

 

자식이 없을 때,1개일 때,2개일 때를 모두 따져줘야 했습니다..

 

(코드가 정말 맘에 안들지만... 나중에 리팩터링해서 수정하도록 하겠습니다.)

 

    public mutating func remove(value:T) {
        if root == nil {
            return
        }
        
        var parent:BSTNode = root!
        var node:BSTNode? = root
        
        while node != nil {
            if node!.value == value {
                //자식이 2개 있을 때
                if node?.leftChild != nil && node?.rightChild != nil {
                    let rightChild = node?.rightChild
                    node = findLeftChildMaxValue(node: node!.leftChild!)
                    node?.rightChild = rightChild
                //자식이 왼쪽만 있을 떄
                }else if node?.leftChild != nil{
                    if node!.value < parent.value {
                        parent.leftChild = node?.leftChild
                    }else {
                        parent.rightChild = node?.leftChild
                    }
                //자식이 오른쪽만 있을 때
                }else if node?.rightChild != nil {
                    if node!.value < parent.value {
                        parent.leftChild = node?.rightChild
                    }else {
                        parent.rightChild = node?.rightChild
                    }
                //자식이 없을 때
                }else {
                    if node?.value == root?.value {
                        root = nil
                        return
                    }
                    if node!.value < parent.value {
                        parent.leftChild = nil
                    }else {
                        parent.rightChild = nil
                    }
                }
                if value == root?.value {
                    self.root = node!
                }
                break
            }
            parent = node!
            node = node!.value > value ? node?.leftChild : node?.rightChild
        }
    }

 

여기서 만약 자식이 2개라면 오른쪽 자식의 가장 작은 값 또는 왼쪽 자식의 가장 큰 값을 대체해야 하기 때문에

 

아래는 오른쪽 자식의 가장 작은 값을 구하는 메서드이고,

 

private func findRightChildMinValue(node:BSTNode<T>) -> BSTNode<T> {
        var parent:BSTNode<T> = node
        var leftChild:BSTNode? = node.leftChild
        
        while leftChild?.leftChild != nil {
            leftChild = leftChild?.leftChild!
            parent = leftChild!
        }
        if leftChild?.rightChild != nil {
            parent.rightChild = leftChild?.rightChild
        }
        return leftChild ?? parent
    }

 

아래는 왼쪽 자식의 가장 큰 값을 구하는 메서드 입니다.

 

    private func findLeftChildMaxValue(node:BSTNode<T>) -> BSTNode<T> {
        var parent:BSTNode<T> = node
        var rightChild:BSTNode? = node.rightChild
        while rightChild?.rightChild != nil {
            rightChild = rightChild?.rightChild!
            parent = rightChild!
        }
        if rightChild?.leftChild != nil {
            parent.leftChild = rightChild?.leftChild
        }
        return rightChild ?? parent
    }

 

삭제 또한 Insert에서 넣은 값으로 테스트 해보겠습니다.

 

1. 자식이 없을 경우

 

먼저 자식이 없을 때의 값인 3을 삭제해 보겠습니다.

 

아래와 같이 정상적으로 3이 삭제된 것을 볼 수 있습니다.

 

        bst.remove(value: 3)
        printChild(node: bst.root)
        
//        10의 왼쪽 자식 :Optional(8)
//        10의 오른쪽 자식 :Optional(11)
//        8의 왼쪽 자식 :Optional(1)
//        8의 오른쪽 자식 :nil
//        1의 왼쪽 자식 :nil
//        1의 오른쪽 자식 :nil
//        11의 왼쪽 자식 :nil
//        11의 오른쪽 자식 :nil
        
        //        10
        //    8       11
        //1

 

2. 자식이 1개인 경우

 

그 다음으로 자식이 1개인 경우인 8을 삭제해보겠습니다.

 

아래와 같이 8이 삭제되고 그 아래 자식으로 대체된 것을 볼 수 있습니다.

 

        bst.remove(value: 8)
        printChild(node: bst.root)
        
//        10의 왼쪽 자식 :Optional(1)
//        10의 오른쪽 자식 :Optional(11)
//        1의 왼쪽 자식 :nil
//        1의 오른쪽 자식 :Optional(3)
//        3의 왼쪽 자식 :nil
//        3의 오른쪽 자식 :nil
//        11의 왼쪽 자식 :nil
//        11의 오른쪽 자식 :nil
        
        //        10
        //    1       11
        //       3

 

3. 자식이 2개인 경우

 

그 다음으로 자식이 2개인 경우인 10을 삭제해보겠습니다.

 

(저는 왼쪽 자식 중 가장 큰 값으로 대체하는 메서드를 사용했습니다.)

 

아래와 같이 10이 삭제되고 왼쪽 값 중 가장 큰 값인 8로 대체되는 것을 볼 수 있습니다.

 

        bst.remove(value: 10)
        printChild(node: bst.root)
        
//        8의 왼쪽 자식 :Optional(1)
//        8의 오른쪽 자식 :Optional(11)
//        1의 왼쪽 자식 :nil
//        1의 오른쪽 자식 :Optional(3)
//        3의 왼쪽 자식 :nil
//        3의 오른쪽 자식 :nil
//        11의 왼쪽 자식 :nil
//        11의 오른쪽 자식 :nil
        
        //        8
        //    1       11
        //       3