优先队列Pq

 

# 基于二叉堆实现

local arrayext = require("arrayext")

local Pq = {}
Pq.__cname = "util.Pq"
Pq.__index = Pq

local ClearTypeEnum = {
    Reset = nil,
    FillNil = 1,
    Keep = 2,
}
Pq.ClearTypeEnum = ClearTypeEnum

--[[
       1
     /   \
    2     3
   / \   / \
  4   5 6   7
 /
8
这边使用数组来保存二叉树
#通过父节点获取子节点: 比如: 2的2个子节点: 2*2和2*2+1, 3的2个子节点3*2和3*2+1
#通过子节点获取父节点: 5的父节点floor(5/2)=2, 6的父节点floor(6/2)=3
]]

function Pq.new(cmpFunc)
    local obj = {}
    setmetatable(obj, Pq)
    obj:ctor(cmpFunc)
    return obj
end

function Pq:ctor(cmpFunc)
    self.cmpFunc = cmpFunc
    self.arr = {}
    self.count = 0
end

function Pq:Clear(clearType)
    if clearType == ClearTypeEnum.Reset then
        self.arr = {}
    elseif clearType == ClearTypeEnum.FillNil then
        for i=1,self.count do
            self.arr[i] = nil
        end
    end
    self.count = 0
end

function Pq:GetCount()
    return self.count
end

function Pq:Push(v)
    self.count = self.count + 1
    self.arr[self.count] = v

    local arr = self.arr
    local cmpFunc = self.cmpFunc

    local temp = self.count
    while temp > 1 do
        local parent = math.floor(temp / 2)
        if cmpFunc(arr[temp], arr[parent]) then
            arrayext.Swap(arr, temp, parent)
            temp = parent
        else
            break
        end
    end
end

local function _UpdateHeap(arr, cmpFunc, count)
    local temp = 1
    while true do
        local lChild = temp * 2
        if lChild > count then break end

        local rChild = lChild + 1
        if rChild > count then
            if not cmpFunc(arr[temp], arr[lChild]) then
                arrayext.Swap(arr, temp, lChild)
                temp = lChild
            else
                break
            end
        else
            local lValue = arr[lChild]
            local rValue = arr[rChild]
            if cmpFunc(lValue, rValue) then --最小堆, 小的与parent比; 最大堆, 大的与parent比;
                if not cmpFunc(arr[temp], arr[lChild]) then
                    arrayext.Swap(arr, temp, lChild)
                    temp = lChild
                else
                    break
                end
            else
                if not cmpFunc(arr[temp], arr[rChild]) then
                    arrayext.Swap(arr, temp, rChild)
                    temp = rChild
                else
                    break
                end
            end
        end
    end
end

function Pq:Pop()
    local count = self.count
    if count < 1 then return nil end
    local newCount = count - 1

    local v = self.arr[1]
    self.arr[1] = nil
    --和最后一个元素交换, 然后再重新构建二叉堆
    if count > 1 then
        arrayext.Swap(self.arr, count, 1)
        _UpdateHeap(self.arr, self.cmpFunc, newCount)
    end
    self.count = newCount

    return v
end

function Pq:Peek()
    if self.count < 1 then return nil end
    return self.arr[1]
end

function Pq:__tostring()
    local count = self.count
    if count < 1 then return "" end

    local strTb = {}
    local arr = arrayext.NewFrom(self.arr)

    while count > 0 do
        local newCount = count - 1

        local v = arr[1]
        table.insert(strTb, tostring(v))
        if count > 1 then
            arrayext.Swap(arr, count, 1)
            _UpdateHeap(arr, self.cmpFunc, newCount)
        end
        count = newCount
    end

    return table.concat(strTb, ",")
end

function Pq:GetIterator()
    if nil == self.iterator then
        self.iterator = function(tb, key)
            return self:Pop()
        end
    end
    return self.iterator
end

return Pq

 

【参考】

二叉堆 - 简书 (jianshu.com)

浅析基础数据结构 - 二叉堆 - 知乎 (zhihu.com)