四十三庵

蔀の雑記帳

Heap Sort(ヒープソート)

データ構造とアルゴリズム (新・情報 通信システム工学)

  • Heapを使ってSortする
  • Performance of heap sort is O(n log n) in best, worst, and average case

実装

Heapの実装にどっかバグがあって、上手くソートできなかった。
一日格闘してダメだったので、SACの模範解答見てお茶を濁した。

import Foundation

struct Heap<Element> {
    var elements = [Element]()
    let priorityFunction: (Element, Element) -> Bool

    init(elements: [Element], priorityFunction: @escaping (Element, Element) -> Bool) {
        self.priorityFunction = priorityFunction
        elements.forEach { insert(element: $0); print(self.elements) }
    }

    // 末尾に要素を挿入し、親要素と比較して、大小関係によって入れ替える
    mutating func insert(element: Element) {
        elements.append(element)
        print(elements)
        var childIndex = lastIndex()
        while rootIndex() < childIndex {
            guard let parentIndex = parentIndex(index: childIndex) else { return }
            if priorityFunction(elements[childIndex], elements[parentIndex]) {
                elements.swapAt(childIndex, parentIndex)
                childIndex = parentIndex
            } else {
                return
            }
        }
    }

    // rootの要素を削除し、末尾の要素をrootにし、左右の子と大小関係を比較して入れ替える
    mutating func delete() {
        var index = rootIndex()
        elements.swapAt(index, lastIndex())
        elements.removeLast()
        while index < lastIndex() {
            guard let leftIndex = leftChildIndex(index: index) else { return } // 左の子要素なし
            if priorityFunction(elements[leftIndex], elements[index]) {
                guard let rightIndex = rightChildIndex(index: index) else { return } // 右の子要素なし
                if priorityFunction(elements[rightIndex], elements[leftIndex]) {
                    elements.swapAt(index, rightIndex)
                    index = rightIndex
                } else {
                    elements.swapAt(index, leftIndex)
                    index = leftIndex
                }
            }
            guard let rightIndex = rightChildIndex(index: index) else { break } // 右の子要素なし(二回やってるが致し方なし)
            if priorityFunction(elements[rightIndex], elements[index]) == true {
                elements.swapAt(index, rightIndex)
                index = rightIndex
            }
            return
        }
    }

    func rootIndex() -> Int {
        return 0
    }
    
    func lastIndex() -> Int {
        return elements.indices.last!
    }

    func parentIndex(index i: Int) -> Int? {
        guard rootIndex() < i else { return nil }
        return (i - 1) / 2
    }

    func leftChildIndex(index i: Int) -> Int? {
        guard (i * 2) + 1 <= lastIndex() else { return nil }
        return (i * 2) + 1
    }

    func rightChildIndex(index i: Int) -> Int? {
        guard (i * 2) + 2 <= lastIndex() else { return nil }
        return (i * 2) + 2
    }
}

extension Heap {
    mutating func sort() -> [Element] {
        for i in stride(from: lastIndex(), through: 1, by: -1) {
            elements.swapAt(0, i)
            pushDown(firstIndex: 0, lastIndex: i - 1)
        }
        return elements
    }
    
    mutating func pushDown(firstIndex: Int, lastIndex: Int) {
        var index = firstIndex
        while index < (lastIndex - 1) / 2 {
            guard let leftIndex = leftChildIndex(index: index) else { return }
            var smallerChildIndex = leftIndex
            if let rightIndex = rightChildIndex(index: index) {
                if priorityFunction(elements[rightIndex], elements[leftIndex]) { smallerChildIndex = rightIndex }
            }
            if priorityFunction(elements[smallerChildIndex], elements[index]) {
                elements.swapAt(index, smallerChildIndex)
                index = smallerChildIndex
            } else {
                break
            }
        }
    }
}

func maxHeapSort<Element: Comparable>(elements: [Element]) -> [Element] {
    var heap = Heap(elements: elements, priorityFunction: <)
    return heap.sort()
}

func minHeapSort<Element: Comparable>(elements: [Element]) -> Heap<Int> {
    var heap = Heap(elements: elements, priorityFunction: >)
    return heap as! Heap<Int>
//    return heap.sort()
}

//var array = [1, 2, 3, 4, 5]
var array = (1...1000).map { $0 }
array[0] = 1001
array[1] = 1002
//order = 0
print(maxHeapSort(elements: array))
//print("計算量はO(", order, ")")

//order = 0
//print(minHeapSort(elements: array))
//print("計算量はO(", order, ")")