Skip to content

Instantly share code, notes, and snippets.

@leegao
Created July 10, 2011 15:56
Show Gist options
  • Save leegao/1074642 to your computer and use it in GitHub Desktop.
Save leegao/1074642 to your computer and use it in GitHub Desktop.
Heap based Priority Queue for 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
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
@devurandom
Copy link

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

@lwt1104
Copy link

lwt1104 commented Aug 5, 2014

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.

@gtfcugb
Copy link

gtfcugb commented Feb 13, 2015

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

@AKST
Copy link

AKST commented Aug 15, 2016

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment