스위프트로 구현하는 자료구조 5: 힙(Heap)

힙 (Heap)

힙은 이진트리를 사용해서 데이터들을 반정렬 상태로 유지하는 자료구조입니다. 힙은 트리의 루트 노드에 데이터들의 최솟값 혹은 최대값을 저장하는데요, 최솟값을 루트에 둘 때는 최소힙, 최대값을 루트에 두면 최대힙이라고 합니다.

완전 이진 트리

이진 트리는 트리를 구성하는 모든 노드가 최대 두 개의 자식 노드를 가지는 자료구조입니다. 그리고 완전 이진 트리는 각 높이에 있는 모든 노드가 자식 노드를 가져야 다음 높이에 새 노드를 추가할 수 있고, 가장 왼쪽부터 차례대로 노드를 채워야합니다.

이런 이진 트리는 배열로도 구현이 가능한데요, 루트 노드를 인덱스 1로 삼으면 어떤 index 번째 노드의 왼쪽 자식 노드는 index * 2, 오른쪽 자식 노드는 (index * 2) + 1 로 구할 수 있습니다.

우리는 힙을 이 배열 기반의 이진 트리를 통해서 구현할거에요!

그리고 힙의 이진 트리는 트리를 구성하는 각 서브 트리들도 모두 힙을 만족해야합니다. 

 

뼈대 만들기

import Foundation
 
struct Heap<T: Comparable> {
    private var elements: [T] = []
    private let sortFunction: (T, T) -> Bool
    
    var isEmpty: Bool {
        return self.elements.count == 1
    }
    var peek: T? {
        if self.isEmpty { return nil }
        return self.elements[1]
    }
    var count: Int {
        return self.elements.count - 1
    }
    
    init(elements: [T] = [], sortFunction: @escaping (T, T) -> Bool) {
        if !elements.isEmpty {
            self.elements = [elements.first!] + elements
        } else {
            self.elements = elements
        }
        self.sortFunction = sortFunction
        if elements.count > 1 {
            self.buildHeap()
        }
    }
    
    func leftChild(of index: Int) -> Int {
        return index * 2 
    }
    func rightChild(of index: Int) -> Int {
        return index * 2 + 1
    }
    func parent(of index: Int) -> Int {
        return (index) / 2
    }
 }

먼저 기본적인 프로퍼티들과 생성자, 메서드를 구현해주었습니다. 

 

elements 는 실제로 데이터들을 저장할 이진트리로 표현될 배열이고, sortFunction은 최소힙, 최대힙의 기준점으로 삼을 정렬 로직을 가진 클로저입니다. 만약 저장되는 데이터가 정수이고, sortFunction이 오름차순 정렬이면 최소힙이 되고, 내림차순 정렬이면 최대힙이 되겠죠?

 

isEmpty 프로퍼티는 현재 이진트리의 요소들이 1개인지 확인하는 것으로 판별합니다. 배열의 0번째 인덱스는 사용하지 않는 값이니 배열에 값이 하나만 있을 때 트리가 비워져 있는 것입니다.

 

leftChild, rightChild, parent는 위에서 설명했던대로 부모와 자식노드의 인덱스를 구하는 메서드이구요, 생성자의 buildHeap는 잠시 뒤에 설명하겠습니다!

 

삽입(insert) - O(log N)

힙의 삽입은 다음과 같은 과정으로 진행됩니다.

이렇게 최대힙이 있을 때 새로운 루트노드가 될 노드를 하나 삽입한다고 해보겠습니다.

먼저 가장 왼쪽 리프노드에 새로 추가할 노드를 붙여줍니다. 

 

이제 최대 힙을 유지하기 위해 트리를 재조정해주어야 하는데요, 가장 끝에 노드를 붙였으니 이 노드를 올라갈 수 있는 최대한 올려주어야 합니다. 그래야 트리의 각 서브 트리들도 최대 힙을 만족하니까요.

먼저 9는 부모 노드인 1보다 크기 때문에 두 노드를 교환해줍니다. 

9의 새로운 부모 노드 4도 9보다 작네요. 다시 교환합시다.

그리고 다시 새로운 부모노드인 7보다도 크기 때문에 새로운 루트노드는 9가 됩니다.

전체 트리와 모든 서브트리가 최대힙을 만족하는 것을 확인할 수 있습니다!

 

삽입 구현 - swimUp

결국 삽입 연산은 

 

1) 트리 제일 끝에 새 노드 붙이기

2) 새 노드를 올릴 수 있는데까지 올리기

 

두 단계로 진행됩니다.

 

이때 2번 과정을 저에게 자료구조를 가르쳐주신 교수님은 swimUp 이라고 표현하셨는데요, 마치 수영해서 올라가는 것 같다고...

 

swimUp을 구현해봅시다.

    mutating func swimUp(from index: Int) {
        var index = index
        while index != 1 && self.sortFunction(self.elements[index], self.elements[self.parent(of: index)]) {
            self.elements.swapAt(index, self.parent(of: index))
            index = self.parent(of: index)
        }
    }

끝났습니다..ㅋㅋㅋㅋ

 

루트노드에 다다르거나, 현재 노드가 부모 노드보다 우선순위가 낮아지기 전까지 노드와 부모노드를 계속 교환해줍니다. 위  그림의 예시에 대한 알고리즘 동작 과정을 써보면 

 

1. 현재 인덱스 8부터 swimUp 시작!

2. 현재 인덱스의 값 9는 부모 인덱스 4의 값 1보다 크기 때문에 두 노드의 값을 교환하고 현재 인덱스를 부모 인덱스인 4로 변경합니다.

3. 현재 인덱스 4의 값 9는 부모 인덱스 2의 값 4보다 크기 때문에 두 노드의 값을 교환하고 현재 인덱스를 부모 인덱스인 2로 변경합니다.

4. 현재 인덱스 2의 값 9는 부모 인덱스 1의 값 7보다 크기 때문에 두 노드의 값을 교환하고 현재 인덱스를 부모 인덱스인 1로 변경합니다.

5. 현재 인덱스는 1이기 때문에 while문을 빠져나와서 함수는 종료됩니다.

 

이 알고리즘은 최악의 경우에는 지금처럼 리프노드부터 루트노드까지 거슬러 올라가며 교환을 진행하기 때문에 시간복잡도는 트리의 높이에 대한 노드 개수인 O(log N)가 됩니다.

 

따라서 전체 삽입 과정은 트리 끝에 데이터를 하나 삽입하고 swimUp을 수행하는 O(1) + O(log N) = O(log N)가 됩니다.

    mutating func insert(node: T) {
        if self.elements.isEmpty {
            self.elements.append(node)
            return
        }
        self.elements.append(node)
        self.swimUp(from: self.elements.endIndex - 1)
    }

 

삭제(delete) - O(log N)

삭제는 루트노드에 있는 가장 우선순위가 높은 노드를 꺼내고 트리를 재조정 하는 작업이 필요합니다.

여기서 삭제를 하면 루트노드가 나오겠죠? 힙에서의 삭제는 루트노드와 마지막 리프 노드를 교환해주는 것으로 시작됩니다.

이렇게 교환되면 루트노드 였던 노드를 제거해주는 것은 간단하겠죠? 트리에서 연결을 끊어줍니다. 

우리는 이진 트리가 배열로 구현되어 있기 때문에 노드의 삭제가 어렵지 않습니다. 그냥 배열에서 마지막 값을 빼주면 되니까요!

 

대신 최대힙이 깨졌기 때문에 다시 트리를 재조정해야 합니다.

루트 노드의 자식 노드들 중에 우선 순위가 더 높은 노드를 골라 현재 루트 노드와 변경해줍시다.

 

이제 원래 루트노드에 있던 노드가 이동하면서 서브트리의 새로운 루트 노드가 됩니다. 따라서 서브트리가 최대 힙을 만족하는지 확인해봐야겠죠? 이번에도 자식 노드들 중에 우선순위가 가장 높은 노드와 서브트리의 루트노드를 비교해봅니다. 

 

루트노드는 5이고, 자식노드는 하나 뿐인 2이기 때문에 현재 트리의 상태가 유지됩니다.

이렇게 하면 모든 서브 트리의 최대 힙을 유지하면서 루트노드를 제거할 수 있게됩니다.

 

삭제 구현 - diveDown

위에 설명한대로 삭제 연산은 루트 노드와 리프 노드의 교환, 그리고 트리의 재조정의 순서로 진행됩니다.

 

swimUp을 명명하신 교수님께서는 루트노드에서부터 하위 노드들을 갱신하면서 점진적으로 내려가는 함수를 diveDown으로 명명하셨는데요. 이것도 위에서부터 다이빙해서 내려오는 것 같다고...

    mutating func diveDown(from index: Int) {
        var higherPriority = index
        let leftIndex = self.leftChild(of: index)
        let rightIndex = self.rightChild(of: index)
        
        // 왼쪽 자식 노드가 더 우선순위가 높을 때
        if leftIndex < self.elements.endIndex && self.sortFunction(self.elements[leftIndex], self.elements[higherPriority]) {
            higherPriority = leftIndex
        }
        // 오른쪽 자식 노드가 더 우선순위가 높을 때
        if rightIndex < self.elements.endIndex && self.sortFunction(self.elements[rightIndex], self.elements[higherPriority]) {
            higherPriority = rightIndex
        }
        // 교환이 필요한 경우는 교환 후 서브트리로 이동
        if higherPriority != index {
            self.elements.swapAt(higherPriority, index)
            self.diveDown(from: higherPriority)
        }
    }

diveDown의 구현은 이렇게 됩니다.

 

더 이상 교환이 필요없을 때까지 왼쪽 자식노드와 오른쪽 자식 노드를 현재 루트 노드와 비교하면서 우선순위가 더 높은 노드와 교환을 진행해주고, 교환이 되면 새로운 노드를 루트 노드로 삼아서 diveDown을 재귀적으로 진행합니다.

 

이번에도 루트와 리프를 교환하는 시간복잡도 O(1), 최악의 경우 루트노드부터 리프노드까지 교환을 진행하는 시간복잡도 O(log N)로 삭제 연산은 O(log N)의 시간복잡도가 요구됩니다.

 

배열을 힙으로 - Build Heap

    init(elements: [T] = [], sortFunction: @escaping (T, T) -> Bool) {
        if !elements.isEmpty {
            self.elements = [elements.first!] + elements
        } else {
            self.elements = elements
        }
        self.sortFunction = sortFunction
        if elements.count > 1 {
            self.buildHeap()
        }
    }

아까 생성자를 이렇게 구현했었는데요, 만약 생성자로 받은 배열에 값이 들어 있다면 배열 전체를 힙으로 다시 재구성 해주어야하기 때문에 이를 위해서 buildHeap이라는 메서드를 호출하고 있습니다.

 

    mutating func buildHeap() {
        for index in (1...(self.elements.count / 2)).reversed() {
            self.diveDown(from: index)
        }
    }

buildHeap은 이렇게 구현되어 있습니다. 배열에 들어있는 요소의 개수를 2로 나누면 마지막 리프노드의 부모 노드가 나오기 때문에 이 노드가 리프노드를 제외하고는 가장 마지막 노드가 됩니다.

이 노드부터 루트노드까지 역순으로 각 서브트리의 루트노드들을 순회하면서 diveDown을 호출하면 가장 마지막 서브트리부터 순서대로 힙을 재구성하게 됩니다.

 

전체 코드

	//
	// Heap.swift
	// Heap
	//
	// Created by 전여훈 on 2021/12/12.
	//
	import Foundation
	struct Heap<T: Comparable> {
	private var elements: [T] = []
	private let sortFunction: (T, T) -> Bool
	var isEmpty: Bool {
	return self.elements.count == 1
	}
	var peek: T? {
	if self.isEmpty { return nil }
	return self.elements.last!
	}
	var count: Int {
	return self.elements.count - 1
	}
	init(elements: [T] = [], sortFunction: @escaping (T, T) -> Bool) {
	if !elements.isEmpty {
	self.elements = [elements.first!] + elements
	} else {
	self.elements = elements
	}
	self.sortFunction = sortFunction
	if elements.count > 1 {
	self.buildHeap()
	}
	}
	func leftChild(of index: Int) -> Int {
	return index * 2
	}
	func rightChild(of index: Int) -> Int {
	return index * 2 + 1
	}
	func parent(of index: Int) -> Int {
	return (index) / 2
	}
	mutating func add(element: T) {
	self.elements.append(element)
	}
	mutating func diveDown(from index: Int) {
	var higherPriority = index
	let leftIndex = self.leftChild(of: index)
	let rightIndex = self.rightChild(of: index)
	if leftIndex < self.elements.endIndex && self.sortFunction(self.elements[leftIndex], self.elements[higherPriority]) {
	higherPriority = leftIndex
	}
	if rightIndex < self.elements.endIndex && self.sortFunction(self.elements[rightIndex], self.elements[higherPriority]) {
	higherPriority = rightIndex
	}
	if higherPriority != index {
	self.elements.swapAt(higherPriority, index)
	self.diveDown(from: higherPriority)
	}
	}
	mutating func swimUp(from index: Int) {
	var index = index
	while index != 1 && self.sortFunction(self.elements[index], self.elements[self.parent(of: index)]) {
	self.elements.swapAt(index, self.parent(of: index))
	index = self.parent(of: index)
	}
	}
	mutating func buildHeap() {
	for index in (1...(self.elements.count / 2)).reversed() {
	self.diveDown(from: index)
	}
	}
	mutating func insert(node: T) {
	if self.elements.isEmpty {
	self.elements.append(node)
	}
	self.elements.append(node)
	self.swimUp(from: self.elements.endIndex - 1)
	}
	mutating func remove() -> T? {
	if self.isEmpty { return nil }
	self.elements.swapAt(1, elements.endIndex - 1)
	let deleted = elements.removeLast()
	self.diveDown(from: 1)
	return deleted
	}
	}

view raw Heap.swift hosted with ❤ by GitHub

 

 

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

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