-
Notifications
You must be signed in to change notification settings - Fork 9
/
heap.lua
60 lines (53 loc) · 1.31 KB
/
heap.lua
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
-- Localize globals
local math_floor, setmetatable, table_insert = math.floor, setmetatable, table.insert
-- Set environment
local _ENV = {}
setfenv(1, _ENV)
local metatable = {__index = _ENV}
function less_than(a, b) return a < b end
--> empty min heap
function new(less_than)
return setmetatable({less_than = less_than}, metatable)
end
function push(self, value)
table_insert(self, value)
local function heapify(index)
if index == 1 then
return
end
local parent = math_floor(index / 2)
if self.less_than(self[index], self[parent]) then
self[parent], self[index] = self[index], self[parent]
heapify(parent)
end
end
heapify(#self)
end
function pop(self)
local value = self[1]
local last = #self
if last == 1 then
self[1] = nil
return value
end
self[1], self[last] = self[last], nil
last = last - 1
local function heapify(index)
local left_child = index * 2
if left_child > last then
return
end
local smallest_child = left_child + 1
if smallest_child > last or self.less_than(self[left_child], self[smallest_child]) then
smallest_child = left_child
end
if self.less_than(self[smallest_child], self[index]) then
self[index], self[smallest_child] = self[smallest_child], self[index]
heapify(smallest_child)
end
end
heapify(1)
return value
end
-- Export environment
return _ENV