Comparison

util/indexedbheap.lua @ 5876:f62ad84811df

util.indexedbheap: A priority queue implementation with a reverse index with no per-entry memory allocation.
author Waqas Hussain <waqas20@gmail.com>
date Wed, 30 Oct 2013 17:30:35 -0400
child 6151:a34a14054532
comparison
equal deleted inserted replaced
5875:e4150e6720f7 5876:f62ad84811df
1
2 local setmetatable = setmetatable;
3 local math_floor = math.floor;
4 local t_remove = table.remove;
5
6 local function _heap_insert(self, item, sync, item2, index)
7 local pos = #self + 1;
8 while true do
9 local half_pos = math_floor(pos / 2);
10 if half_pos == 0 or item > self[half_pos] then break; end
11 self[pos] = self[half_pos];
12 sync[pos] = sync[half_pos];
13 index[sync[pos]] = pos;
14 pos = half_pos;
15 end
16 self[pos] = item;
17 sync[pos] = item2;
18 index[item2] = pos;
19 end
20
21 local function _percolate_up(self, k, sync, index)
22 local tmp = self[k];
23 local tmp_sync = sync[k];
24 while k ~= 1 do
25 local parent = math_floor(k/2);
26 if tmp < self[parent] then break; end
27 self[k] = self[parent];
28 sync[k] = sync[parent];
29 index[sync[k]] = k;
30 k = parent;
31 end
32 self[k] = tmp;
33 sync[k] = tmp_sync;
34 index[tmp_sync] = k;
35 return k;
36 end
37
38 local function _percolate_down(self, k, sync, index)
39 local tmp = self[k];
40 local tmp_sync = sync[k];
41 local size = #self;
42 local child = 2*k;
43 while 2*k <= size do
44 if child ~= size and self[child] > self[child + 1] then
45 child = child + 1;
46 end
47 if tmp > self[child] then
48 self[k] = self[child];
49 sync[k] = sync[child];
50 index[sync[k]] = k;
51 else
52 break;
53 end
54
55 k = child;
56 child = 2*k;
57 end
58 self[k] = tmp;
59 sync[k] = tmp_sync;
60 index[tmp_sync] = k;
61 return k;
62 end
63
64 local function _heap_pop(self, sync, index)
65 local size = #self;
66 if size == 0 then return nil; end
67
68 local result = self[1];
69 local result_sync = sync[1];
70 index[result_sync] = nil;
71 if size == 1 then
72 self[1] = nil;
73 sync[1] = nil;
74 return result, result_sync;
75 end
76 self[1] = t_remove(self);
77 sync[1] = t_remove(sync);
78 index[sync[1]] = 1;
79
80 _percolate_down(self, 1, sync, index);
81
82 return result, result_sync;
83 end
84
85 local indexed_heap = {};
86
87 function indexed_heap:insert(item, priority, id)
88 if id == nil then
89 id = self.current_id;
90 self.current_id = id + 1;
91 end
92 self.items[id] = item;
93 _heap_insert(self.priorities, priority, self.ids, id, self.index);
94 return id;
95 end
96 function indexed_heap:pop()
97 local priority, id = _heap_pop(self.priorities, self.ids, self.index);
98 if id then
99 local item = self.items[id];
100 self.items[id] = nil;
101 return priority, item, id;
102 end
103 end
104 function indexed_heap:peek()
105 return self.priorities[1];
106 end
107 function indexed_heap:reprioritize(id, priority)
108 local k = self.index[id];
109 if k == nil then return; end
110 self.priorities[k] = priority;
111
112 k = _percolate_up(self.priorities, k, self.ids, self.index);
113 k = _percolate_down(self.priorities, k, self.ids, self.index);
114 end
115 function indexed_heap:remove_index(k)
116 local size = #self.priorities;
117
118 local result = self.priorities[k];
119 local result_sync = self.ids[k];
120 local item = self.items[result_sync];
121 if result == nil then return; end
122 self.index[result_sync] = nil;
123 self.items[result_sync] = nil;
124
125 self.priorities[k] = self.priorities[size];
126 self.ids[k] = self.ids[size];
127 self.index[self.ids[k]] = k;
128 t_remove(self.priorities);
129 t_remove(self.ids);
130
131 k = _percolate_up(self.priorities, k, self.ids, self.index);
132 k = _percolate_down(self.priorities, k, self.ids, self.index);
133
134 return result, item, result_sync;
135 end
136 function indexed_heap:remove(id)
137 return self:remove_index(self.index[id]);
138 end
139
140 local mt = { __index = indexed_heap };
141
142 local _M = {
143 create = function()
144 return setmetatable({
145 ids = {}; -- heap of ids, sync'd with priorities
146 items = {}; -- map id->items
147 priorities = {}; -- heap of priorities
148 index = {}; -- map of id->index of id in ids
149 current_id = 1.5
150 }, mt);
151 end
152 };
153 return _M;