Heap based Priority Queue for Lua

example.lua

-- A priority queue that sorts based on euclidean distance

pq = pqueue(function(point1, point2)
    return (point1[1]-point2[1])^2 + (point1[2]-point2[2])^2
end)

for i=1, 20 do
    pq:push{math.random(100), math.random(100)}
end

print(unpack(pq:pop())) -- get the point closest to the origin

pqueue.lua

function pqueue(cmp, initial)
     cmp = cmp or function(a,b) return a<b end
     local pq = setmetatable({}, {
          __index = {
               push = function(self, v)
                    table.insert(self, v)
                    local next = #self
                    local prev = math.floor(next/2)
                    while next > 1 and self[next] < self[prev] do
                         -- swap up
                         self[next], self[prev] = self[prev], self[next]
                         next = prev
                         prev = math.floor(next/2)
                    end
                    self[next] = v
                    self.size = self.size + 1
               end,
               size = 0,
               pop = function(self)
                    if #self < 2 then
                         return table.remove(self)
                    end
                    local r = self[1]
                    self[1] = table.remove(self)
                    local root = 1
                    if #self > 1 then
                         local size = #self
                         local v = self[root]
                         while root < size do
                              local child = 2*root
                              if child < size then
                                   if child+1 < size and self[child+1] < self[child] then child = child + 1 end
                                   if cmp(self[child], self[root]) then
                                        -- swap down
                                        self[root], self[child] = self[child], self[root]
                                        root = child
                                   else
                                        self[root] = v
                                        break
                                   end
                              else
                                   self[root] = v
                                   break
                              end
                         end
                    end
                    self.size = self.size - 1
                    return r
               end}})
     for _,el in ipairs(initial or {}) do pq:push(el) end
     return pq
end

Is there a reason line 9 doesn't also use cmp?

@Quantumplation: You are right. Line 32 has the same issue.

And at least in Lua 5.2 the example is flawed. The comparison function needs to be something like:

function(a,b) return a[1]^2+a[2]^2 < b[1]^2+b[2]^2 end

It is quite flawed indeed. I recommend not to use it -- try this one instead, I just reimplemented it:

local table = require "table"
local insert = table.insert
local remove = table.remove

local priority_queue = {}

function priority_queue.new(cmp, initial)
    local cmp = cmp or function(a,b) return a < b end

    local pq = setmetatable({}, {
        __index = {
            size = 0,
            push = function(self, v)
                insert(self, v)
                local next = #self
                local prev = (next-next%2)/2
                while next > 1 and cmp(self[next], self[prev]) do
                    self[next], self[prev] = self[prev], self[next]
                    next = prev
                    prev = (next-next%2)/2
                end
            end,
            pop = function(self)
                if #self < 2 then
                    return remove(self)
                end
                local root = 1
                local r = self[root]
                self[root] = remove(self)
                local size = #self
                if size > 1 then
                    local child = 2*root
                    while child <= size do
                        if cmp(self[child], self[root]) then
                            self[root], self[child] = self[child], self[root]
                            root = child
                        elseif child+1 <= size and cmp(self[child+1], self[root]) then
                            self[root], self[child+1] = self[child+1], self[root]
                            root = child+1
                        else
                            break
                        end
                        child = 2*root
                    end
                end
                return r
            end,
            peek = function(self)
                return self[1]
            end,
        }
    })

    for _,el in ipairs(initial or {}) do
        pq:push(el)
    end

    return pq
end

return priority_queue

if cmp(self[child], self[root]) then
self[root], self[child] = self[child], self[root]
root = child
elseif child+1 <= size and cmp(self[child+1], self[root]) then
self[root], self[child+1] = self[child+1], self[root]
root = child+1
else
break
end

This part may be wrong because we should exchange root element with the largest child.

yes,it's wrong ,follow code is ok

                while child <= size do

                    local aBool =   cmp(self[root],self[child]);
                    local bBool =   true;
                    local cBool =   true;
                    if child+1 <= size then
                        bBool =   cmp( self[root],self[child+1]);
                        cBool =   cmp(self[child], self[child+1]);
                    end
                    if aBool and bBool then
                        break;
                    elseif cBool then
                        self[root], self[child] = self[child], self[root]
                        root = child
                    else
                        self[root], self[child+1] = self[child+1], self[root]
                        root = child+1
                    end
                    child = 2*root
                end

@leegao would you consider updating your solution to reflect some the suggestions in this discussion :O)