[Data Structure] AVL 트리 구현해보기 (feat. Swift)

 

안녕하세요 Foma 💻 입니다!

 

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

 

(혹시 안보신 분들은 여기 에서 확인해주세요!)

 

오늘은 AVL 트리를 직접 구현해보려고 합니다.

 

바로 시작할게요~

---

AVLNode

 

기본적으로 이진탐색트리의 노드와 똑같지만 왼쪽과,오른쪽 자식의 높이와 밸런스 팩터를 측정할 변수를 따로

 

만들어었습니다.

 

class AVLNode<T:Comparable> {
    var value:T
    var leftChild:AVLNode?
    var rightChild:AVLNode?
    var height:Int = 0

    var leftHeight:Int {
        return leftChild?.height ?? -1
    }
    
    var rightHeight:Int {
        return rightChild?.height ?? -1
    }
    
    var balanceFactor:Int {
        return leftHeight - rightHeight
    }
    
    init(value:T) {
        self.value = value
    }
}

extension AVLNode {
    var min: AVLNode {
        return leftChild?.min ?? self
    }
}

---

Balance

 

균형인지 아닌지를 체크하는 메서드를 만들어 줍니다.

 

절대값이 2인지 아닌지를 체크해줍니다.

 

 private func isBalanced(_ node: AVLNode<T>) -> Bool {
        return abs(node.balanceFactor) != 2
    }

 

만약 균형이 잡히지 않았다면 균형을 잡아줘야겠죠?

 

밸런스팩터가 2라면 왼쪽으로 치우쳐진 것이고 -2라면 오른쪽으로 치우쳐진 것이므로

 

밸런스 팩터에 맞게 균형을 맞춰줍니다.

 

 private func makeBalanced(_ node: AVLNode<T>) -> AVLNode<T> {
        if node.balanceFactor == 2 {
            return leftBalanced(node)
        }else {
            return rightBalanced(node)
        }
    }

 

 

왼쪽으로 치우쳐 있는 경우 왼쪽 자식의 밸런스 팩터가 -1이라면 오른쪽에 자식이 있는 것이므로 

 

왼쪽으로 회전을 하고 오른쪽으로 회전을 해서 균형을 맞춰줍니다.

 

그게 아니라면 그냥 오른쪽으로 회전해서 균형을 맞춰줍니다.

 

  private func leftBalanced(_ node: AVLNode<T>) -> AVLNode<T> {
        if let leftChild = node.leftChild, leftChild.balanceFactor == -1 {
            return leftRightRotate(node)
        }else {
            return rightRotate(node)
        }
    }

 

오른쪽으로 치우쳐 있는 경우 오른쪽 자식의 밸런스 팩터가 1이라면 왼쪽에 자식이 있는 것이므로

 

오른쪽으로 회전을 하고 왼쪽으로 회전을 해서 균형을 맞춰줍니다.

 

그게 아니라면 그냥 왼쪽으로 회전해서 균형을 맞춰줍니다.

 

  private func rightBalanced(_ node: AVLNode<T>) -> AVLNode<T> {
        if let rightChild = node.rightChild, rightChild.balanceFactor == 1 {
            return rightLeftRotate(node)
        }else {
            return leftRotate(node)
        }
    }

---

Rotate

Right

 

왼쪽으로 치우쳐져 있는 경우는 아래와 같이 오른쪽으로 회전해서 균형을 맞춰줘야 합니다.

 

 

 

 

회전해주는 방법은 가장 첫 번째 노드를 중앙값의 오른쪽 자식으로 만듭니다.

 

중앙값은 두 번째 값이므로 첫 번째 노드의 왼쪽 자식이 되겠죠?

 

균형을 맞춰줬으니 노드의 높이를 조정해줍니다.

 

  private func rightRotate(_ node: AVLNode<T>) -> AVLNode<T> {
        let middle = node.leftChild!
        node.leftChild = middle.rightChild
        middle.rightChild = node
        node.height = max(node.leftHeight, node.rightHeight) + 1
        middle.height = max(middle.leftHeight, middle.rightHeight) + 1
        return middle
    }

 

Left

 

오른쪽으로 치우쳐 있는 경우는 아래와 같이 왼쪽으로 회전해서 균형을 맞춰줘야 합니다.

 

 

 

회전해주는 방법은 가장 첫 번째 노드를 중앙값의 왼쪽 자식으로 만듭니다.

 

그리고 오른쪽으로 회전했을 때와 동일하게 만들어줍니다.

 private func leftRotate(_ node: AVLNode<T>) -> AVLNode<T>{
        let middle = node.rightChild!
        node.rightChild = middle.leftChild
        middle.leftChild = node
        node.height = max(node.leftHeight, node.rightHeight) + 1
        middle.height = max(middle.leftHeight, middle.rightHeight) + 1
        return middle
    }

 

Right-Left

 

아래 첫 번째 그림과 같이 ">" 를 회전해주는 방법은 먼저 두 번째와 세 번째 노드를 오른쪽으로 회전해준 뒤 

 

왼쪽으로 회전해줍니다.

 

 

 

오른쪽 자식을 기준으로 오른쪽으로 회전해준 뒤 왼쪽으로 회전해줍니다.

 

  private func rightLeftRotate(_ node: AVLNode<T>) -> AVLNode<T> {
        node.rightChild = rightRotate(node.rightChild!)
        return leftRotate(node)
    }

 

Left-Right

 

Right-Left와 동일하게 "<" 를 회전해주는 방법은 먼저 두 번째와 세 번째 노드를 왼쪽으로 회전해준 뒤

 

오른쪽으로 회전해줍니다.

 

 

 

  private func leftRightRotate(_ node: AVLNode<T>) -> AVLNode<T> {
        node.leftChild = leftRotate(node.leftChild!)
        return rightRotate(node)
    }

---

Insert

 

삽입을 할 때 재귀를 이용해서 균형이 맞는지 아닌지 확인을 해줍니다.

 

  public mutating func insert(value: T) {
        root = insert(from: root, value: value)
    }
    
    private func insert(from node: AVLNode<T>?, value: T) -> AVLNode<T> {
        guard var node = node else {
            return AVLNode(value: value)
        }
        if value < node.value {
            node.leftChild = insert(from: node.leftChild, value: value)
        } else {
            node.rightChild = insert(from: node.rightChild, value: value)
        }
        if !isBalanced(node) {
            node = makeBalanced(node)
        }
        node.height = max(node.leftHeight, node.rightHeight) + 1
        return node
    }

---

Remove

 

삽입과 동일하게 재귀를 이용해 균형이 맞는지 아닌지 확인합니다.

 

   public mutating func remove(_ value: T) {
        root = remove(node: root, value: value)
    }
    
    private func remove(node: AVLNode<T>?, value: T) -> AVLNode<T>? {
        guard var node = node else {
            return nil
        }
        if value == node.value {
            if node.leftChild == nil && node.rightChild == nil {
                return nil
            }
            if node.leftChild == nil {
                return node.rightChild
            }
            if node.rightChild == nil {
                return node.leftChild
            }
            node.value = node.rightChild!.min.value
            node.rightChild = remove(node: node.rightChild, value: node.value)
        } else if value < node.value {
            node.leftChild = remove(node: node.leftChild, value: value)
        } else {
            node.rightChild = remove(node: node.rightChild, value: value)
        }
        if !isBalanced(node) {
            node = makeBalanced(node)
        }
        node.height = max(node.leftHeight, node.rightHeight) + 1
        return node
    }

---

Test

 

7,4,3,2,5,8,10,11,9 순으로 숫자를 삽입하고 출력을 하면 균형이 잡힌 트리가 나오는 것을 볼 수 있습니다.

 

        var avl = AVLTree<Int>()
        avl.insert(value: 7)
        avl.insert(value: 4)
        avl.insert(value: 3)
        avl.insert(value: 2)
        avl.insert(value: 5)
        avl.insert(value: 8)
        avl.insert(value: 10)
        avl.insert(value: 11)
        avl.insert(value: 9)

//        4
//    3        8
//  2        7    10
//         5     9   11

 

위에서 8을 삭제하고 출력을 해도 정상적으로 균형잡힌 트리가 나오는 것을 볼 수 있습니다.

 

        avl.remove(8)
//        4
//    3        9
//  2        7    10
//         5        11

---

Reference

 

[스위프트 : 자료구조] AVL Tree: 자가 균형 트리: #balance: #트리의 높이: #rotation메소드: #성능오짐