2 local setmetatable = setmetatable;
3 local math_floor = math.floor;
4 local t_remove = table.remove;
6 local function _heap_insert(self, item, sync, item2, index)
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;
21 local function _percolate_up(self, k, sync, index)
23 local tmp_sync = sync[k];
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];
38 local function _percolate_down(self, k, sync, index)
40 local tmp_sync = sync[k];
44 if child ~= size and self[child] > self[child + 1] then
47 if tmp > self[child] then
48 self[k] = self[child];
49 sync[k] = sync[child];
64 local function _heap_pop(self, sync, index)
66 if size == 0 then return nil; end
68 local result = self[1];
69 local result_sync = sync[1];
70 index[result_sync] = nil;
74 return result, result_sync;
76 self[1] = t_remove(self);
77 sync[1] = t_remove(sync);
80 _percolate_down(self, 1, sync, index);
82 return result, result_sync;
85 local indexed_heap = {};
87 function indexed_heap:insert(item, priority, id)
90 self.current_id = id + 1;
92 self.items[id] = item;
93 _heap_insert(self.priorities, priority, self.ids, id, self.index);
96 function indexed_heap:pop()
97 local priority, id = _heap_pop(self.priorities, self.ids, self.index);
99 local item = self.items[id];
100 self.items[id] = nil;
101 return priority, item, id;
104 function indexed_heap:peek()
105 return self.priorities[1];
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;
112 k = _percolate_up(self.priorities, k, self.ids, self.index);
113 k = _percolate_down(self.priorities, k, self.ids, self.index);
115 function indexed_heap:remove_index(k)
116 local result = self.priorities[k];
117 if result == nil then return; end
119 local result_sync = self.ids[k];
120 local item = self.items[result_sync];
121 local size = #self.priorities;
123 self.priorities[k] = self.priorities[size];
124 self.ids[k] = self.ids[size];
125 self.index[self.ids[k]] = k;
127 t_remove(self.priorities);
130 self.index[result_sync] = nil;
131 self.items[result_sync] = nil;
134 k = _percolate_up(self.priorities, k, self.ids, self.index);
135 k = _percolate_down(self.priorities, k, self.ids, self.index);
138 return result, item, result_sync;
140 function indexed_heap:remove(id)
141 return self:remove_index(self.index[id]);
144 local mt = { __index = indexed_heap };
148 return setmetatable({
149 ids = {}; -- heap of ids, sync'd with priorities
150 items = {}; -- map id->items
151 priorities = {}; -- heap of priorities
152 index = {}; -- map of id->index of id in ids