Software / code / prosody
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; |