1. 程式人生 > >演算法-手寫堆

演算法-手寫堆

 

package heap

/**
 * 根節點從1開始
 * parent = k/2
 * lchild = i * 2
 * rchild = i * 2 +1
 */
class MaxHeap {
    var mArr: IntArray
    var mCount: Int = 0

    constructor(size: Int) {
        mArr = IntArray(size + 1)
    }

    fun insert(newVal: Int) {
        assert(mCount + 1 < mArr.size)
        mArr[++mCount] = newVal
        shiftUp(mCount)
    }

    fun extractMax(): Int {
        assert(mCount > 0)
        var result = mArr[1]
        if (mCount > 1) {
            mArr[1] = mArr[mCount]
            shiftDown(1)
        }
        mCount--
        return result
    }

    fun isEmpty(): Boolean {
        return mCount <= 0
    }

    /**
     * 插入元素是插入到末尾,這時候需要與父親對比,如果父親大就交換值
     */
    private fun shiftUp(count: Int) {
        var k = count
        while (k > 1) {
            if (mArr[k / 2] < mArr[k]) {
                var tmp = mArr[k / 2]
                mArr[k / 2] = mArr[k]
                mArr[k] = tmp
            }
            k /= 2
        }
    }

    /**
     * extractMax的時候會先刪除根節點,然後將最後的值補充到根節點,這樣就要
     * 將根節點和孩子對比,然後和最大的孩子交換位置
     */
    private fun shiftDown(start: Int) {
        var k = start//parent
        var j = 2 * k
        while (j <= mCount) {
            //find max from its children
            if (j + 1 <= mCount && mArr[j + 1] > mArr[j]) {
                j++
            }
            //swap(k,j) when k < j
            if (mArr[k] < mArr[j]) {
                var tmp = mArr[k]
                mArr[k] = mArr[j]
                mArr[j] = tmp
            }
            k = j
            j = 2 * k
        }
    }
}

fun main(args: Array<String>) {
    var maxHeap = MaxHeap(100)
    for (i in 1..30) {
        maxHeap.insert((Math.random() * 10).toInt())
    }

    while (!maxHeap.isEmpty()) {
        println(maxHeap.extractMax())
    }

}