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