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
@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