[자료구조] Swift 로 Heap 구현하기

[자료구조] Swift 로 Heap 구현하기

Heap 은 이진트리 형태의 자료구조로 최대값 및 최소값을 빠르게 탐색할수 있도록 설계되었습니다.

 

Heap 에는 2가지 종류가 있는데, 부모노드의 키값이 자식노드의 키값보다 항상 큰 경우는 Max Heap, 부모노드의 키값이 자식노드의 키값보다 항상 작은 경우는 Min Heap 이라고 합니다.

 

그렇기 때문에 Max Heap 의 루트노드에는 항상 최대값이 저장되고 Min Heap 의 루트노드에는 항상 최소값이 저장됩니다.

 

Heap 은 이진트리 형태로 구현되었기 때문에 삽입과 삭제에서 O(logN) 의 시간복잡도를 가집니다.

삽입

MaxHeap 을 예로 들어보겠습니다. Heap 의 삽입은 트리의 가장 바깥쪽 노드에서 시작됩니다.

먼저 삽입할 노드를 트리의 가장 바깥쪽에 붙여줍니다. 하지만 여기서 삽입이 종료된다면 Heap 의 성질을 만족할 수 없기에 노드들의 자리를 조정해줘야 합니다.

 

새로 삽입된 노드 9는 부모노드인 4보다 큰 값을 가지기 때문에 서로 자리를 바꿔줍니다. 바꾼 후에도 부모노드인 8보다 큰 값을 가지기 때문에 한번 더 자리를 바꿔줍니다.

위와 같이 Heap 의 성질을 만족하며 노드의 삽입을 진행할 수 있습니다. 이렇게 새로 삽입된 노드가 부모노드와 비교되며 자리를 찾아 올라가는 과정을 shiftUp 이라고 하겠습니다.

삭제

Heap 의 삭제도 삽입과 마찬가지로 가장 바깥쪽 노드에서 시작합니다.

Heap 은 삭제 시 항상 루트노드를 삭제하고 그 값을 리턴합니다.

 

먼저 루트노드와 가장 바깥쪽 노드의 자리를 바꿔준 다음 원래 루트노드였던 값을 삭제한 후 리턴합니다. 이후에는 루트노드 자리에 Heap 의 성질을 만족하지 않는 노드가 올라와 있을 것입니다.

 

그림에서는 루트노드였던 9와 가장 바깥쪽 노드였던 4의 자리를 바꿔주었습니다. 이후 9를 리턴하고, 임시 루트노드인 4는 자리를 찾아 내려가야 합니다.

 

4의 자식노드인 6과 8중 더 큰 값인 8과 자리를 바꿔줍니다. 이후 자식노드가 되는 3보다는 4의 값이 더 크기 때문에 자리를 바꿔주지 않아도 됩니다.

위와 같이 Heap 의 성질을 만족하며 노드의 삭제를 진행할 수 있습니다. 삭제 후 노드의 자리를 찾아 내려보내는 과정을 shiftDown 이라고 하겠습니다.

코드

	struct Heap<T: Comparable> {
	var heap: [T]
	// maxHeap >, minHeap <
	var compare: (T, T) -> Bool
	var root: T? {
	return heap.isEmpty ? nil : heap[0]
	}
	var isEmpty: Bool {
	return heap.isEmpty
	}
	init(compare: @escaping (T, T) -> Bool) {
	self.heap = []
	self.compare = compare
	}
	// 삽입
	mutating func insert(n: T) {
	heap.append(n)
	shiftUp(i: heap.count - 1)
	}
	// 삭제
	mutating func delete() -> T? {
	if heap.isEmpty {
	return nil
	}
	if heap.count == 1 {
	return heap.removeFirst()
	}
	heap.swapAt(0, heap.count - 1)
	let result = heap.removeLast()
	shiftDown(i: 0)
	return result
	}
	// 삽입 시 노드 상승
	mutating func shiftUp(i: Int) {
	var now = i
	// 현재 노드가 루트 노드보다 하위 노드일때
	while now > 0 {
	// 부모 노드 인덱스
	let parent = (now - 1) / 2
	// 부모 노드와 대소비교, 교환 필요 시
	if compare(heap[now], heap[parent]) {
	heap.swapAt(now, parent)
	// 교환 후 현재 노드를 원래 부모 노드가 있던 자리로 옮겨줌
	now = parent
	} else {
	break
	}
	}
	}
	// 원소 삭제 시 노드 하강
	mutating func shiftDown(i: Int) {
	var now: Int = i
	// 자식 노드 인덱스
	var child: Int = 2 * now + 1
	let count: Int = heap.count
	// 자식 노드가 트리 범위안에 있을때
	while child < count {
	// 왼쪽 자식, 오른쪽 자식 둘다 있을때
	if child + 1 < count {
	child = compare(heap[child], heap[child + 1]) ? child : child + 1
	}
	// 자식 노드와 대소비교 후 교환
	if compare(heap[child], heap[now]) {
	heap.swapAt(now, child)
	now = child // 현재 노드를 자식노드 인덱스로
	child = 2 * now + 1
	} else {
	break
	}
	}
	}
	}

view raw heap.swift hosted with ❤ by GitHub

Heap 에 들어오는 원소들은 대소를 비교하여 자리를 지정해줘야 하기 때문에 타입은 Comparable 해야 합니다.

 

MaxHeap 과 MinHeap 을 선택하여 구현할 수 있도록 인스턴스 생성 시 compare 이라는 인자에 파라미터로 클로져를 입력받습니다.

대소를 비교하는 > < 연산자를 이용하여 각각 MaxHeap 과 MinHeap 을 구현할 수 있습니다.

 

위 코드에서는 Heap 을 배열로 구현하였기 때문에 자식노드의 인덱스는 부모노드의 인덱스에 2를 곱한 값, 혹은 이에 1을 더한 값이 됩니다.

 

앞서 설명한것과 같이 노드를 삽입할 시 shiftUp() 메소드가 실행되고, 노드를 삭제할 시 shiftDown() 메소드가 실행되는 것을 확인할 수 있습니다.