Generic Heap in Swift with Min/Max implementation and tree printing

Heap Swif 3.0

//: [Previous](@previous)

//  Created by Peter Helstrup Jensen on 06-07-2015
//  Copyright (c) 2015 Peter Helstrup Jensen. All rights reserved.
//
//  This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public
//  License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later
//  version.
//  This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied
//  warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
//  You should have received a copy of the GNU General Public License along with this program. If not, see
//  http:/www.gnu.org/licenses/.

import Foundation

protocol HeapType {
    func compare<T: Comparable>(_ first: T, _ second: T) -> Bool
}

class MinHeap: HeapType {
    func compare<T: Comparable>(_ first: T, _ second: T) -> Bool {
        return first < second
    }
}
class MaxHeap: HeapType {
    func compare<T: Comparable>(_ first: T, _ second: T) -> Bool {
        return second < first
    }
}

class Heap<T: Comparable> {
    let type: HeapType
    
    init(heapType: HeapType) {
        self.type = heapType
    }
    
    var data = [T]()
    
    func size() -> Int {
        return data.count
    }
    
    func printHeap() {
        print(data)
    }
    
    func printTree() {
        func tree(_ index: Int, _ depth: inout Int) {
            
            if data.isEmpty { return }
            
            let right = rightIndex(index)
            if right < data.count  {
                depth += 1
                tree(right, &depth)
            }
            
            for _ in 0..<depth*4 {
                print(" ", terminator: "")
            }
            
            print(data[index])
            
            let left = leftIndex(index)
            if left < data.count {
                depth += 1
                tree(left, &depth)
            }
            
            depth -= 1
        }
        
        var d = 0
        tree(0, &d)
    }
    
    
    func percolateUp(_ index: Int) {
        let parent = parentIndex(index)
        if parent != -1 && type.compare(data[index], data[parent]) {
            swap(&data[index], &data[parent])
            percolateUp(parent)
        }
    }
    
    func percolateDown(_ index: Int) {
        let child = childCandidateIndex(index)
        if child != -1 && type.compare(data[child], data[index]) {
            swap(&data[index], &data[child])
            percolateDown(child)
        }
    }
    
    func get() -> T {
        return data[0]
    }
    
    func insert(_ value: T) {
        data.append(value)
        let lastIndex = data.count-1
        percolateUp(lastIndex)
    }
    
    func remove() {
        data[0] = data.last!
        data.removeLast()
        percolateDown(0)
    }
    
    func childCandidateIndex(_ parentIndex: Int) -> Int {
        let rightChildIndex = rightIndex(parentIndex)
        let leftChildIndex = leftIndex(parentIndex)
        
        let rightValue: T? = rightChildIndex < data.count ? data[rightChildIndex] : nil
        let leftValue: T? = leftChildIndex < data.count ? data[leftChildIndex] : nil
        
        if leftValue == nil && rightValue == nil {
            return -1
        }
        else if leftValue != nil && rightValue == nil {
            return leftChildIndex
        }
        else if leftValue == nil && rightValue != nil {
            return rightChildIndex
        }
        else if type.compare(leftValue!, rightValue!) {
            return leftChildIndex
        }
        else {
            return rightChildIndex
        }
    }
    
    func parentIndex(_ childIndex: Int) -> Int {
        return (childIndex-1)/2
    }
    
    func leftIndex(_ parentIndex: Int) -> Int {
        return 2*parentIndex + 1
    }
    
    func rightIndex(_ parentIndex: Int) -> Int {
        return 2*parentIndex + 2
    }
    
    func getSortedArray() -> [T] {
        let heapCopy = Heap(heapType: type)
        heapCopy.data = [T](self.data)
        
        var array = [T]()
        for _ in heapCopy.data {
            array.append(heapCopy.get())
            heapCopy.remove()
        }
        return array
    }
}

//Try changing heap type: MaxHeap() and MinHeap()
let heap = Heap<Int>(heapType: MaxHeap())
heap.insert(31)
heap.insert(32)
heap.insert(1)
heap.insert(15)
heap.insert(17)
heap.insert(12)
heap.insert(4)
heap.insert(35)
heap.insert(79)
heap.insert(44)
heap.insert(34)
heap.insert(54)
heap.insert(2)
heap.insert(9)
heap.insert(454)

print(heap.getSortedArray()) //should be O(n * log(n)) + copy time for convenience testing
print("_____________________________")

for _ in 0..<heap.size() {
    heap.printTree()
    heap.remove()
    print("_____________________________")
}

let sHeap = Heap<String>(heapType: MaxHeap())
sHeap.insert("Abe")
sHeap.insert("Citron")
sHeap.insert("Båd")
sHeap.insert("Dykkerur")
sHeap.insert("Kat")
sHeap.insert("Www")
sHeap.insert("Internet")

for _ in 0..<sHeap.size() {
    sHeap.printTree()
    sHeap.remove()
    print("_____________________________")
}

//: [Next](@next)