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
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
|
local setmetatable = setmetatable;
local math_floor = math.floor;
local t_remove = table.remove;
local function _heap_insert(self, item, sync, item2, index)
local pos = #self + 1;
while true do
local half_pos = math_floor(pos / 2);
if half_pos == 0 or item > self[half_pos] then break; end
self[pos] = self[half_pos];
sync[pos] = sync[half_pos];
index[sync[pos]] = pos;
pos = half_pos;
end
self[pos] = item;
sync[pos] = item2;
index[item2] = pos;
end
local function _percolate_up(self, k, sync, index)
local tmp = self[k];
local tmp_sync = sync[k];
while k ~= 1 do
local parent = math_floor(k/2);
if tmp >= self[parent] then break; end
self[k] = self[parent];
sync[k] = sync[parent];
index[sync[k]] = k;
k = parent;
end
self[k] = tmp;
sync[k] = tmp_sync;
index[tmp_sync] = k;
return k;
end
local function _percolate_down(self, k, sync, index)
local tmp = self[k];
local tmp_sync = sync[k];
local size = #self;
local child = 2*k;
while 2*k <= size do
if child ~= size and self[child] > self[child + 1] then
child = child + 1;
end
if tmp > self[child] then
self[k] = self[child];
sync[k] = sync[child];
index[sync[k]] = k;
else
break;
end
k = child;
child = 2*k;
end
self[k] = tmp;
sync[k] = tmp_sync;
index[tmp_sync] = k;
return k;
end
local function _heap_pop(self, sync, index)
local size = #self;
if size == 0 then return nil; end
local result = self[1];
local result_sync = sync[1];
index[result_sync] = nil;
if size == 1 then
self[1] = nil;
sync[1] = nil;
return result, result_sync;
end
self[1] = t_remove(self);
sync[1] = t_remove(sync);
index[sync[1]] = 1;
_percolate_down(self, 1, sync, index);
return result, result_sync;
end
local indexed_heap = {};
function indexed_heap:insert(item, priority, id)
if id == nil then
id = self.current_id;
self.current_id = id + 1;
end
self.items[id] = item;
_heap_insert(self.priorities, priority, self.ids, id, self.index);
return id;
end
function indexed_heap:pop()
local priority, id = _heap_pop(self.priorities, self.ids, self.index);
if id then
local item = self.items[id];
self.items[id] = nil;
return priority, item, id;
end
end
function indexed_heap:peek()
return self.priorities[1];
end
function indexed_heap:reprioritize(id, priority)
local k = self.index[id];
if k == nil then return; end
self.priorities[k] = priority;
k = _percolate_up(self.priorities, k, self.ids, self.index);
_percolate_down(self.priorities, k, self.ids, self.index);
end
function indexed_heap:remove_index(k)
local result = self.priorities[k];
if result == nil then return; end
local result_sync = self.ids[k];
local item = self.items[result_sync];
local size = #self.priorities;
self.priorities[k] = self.priorities[size];
self.ids[k] = self.ids[size];
self.index[self.ids[k]] = k;
t_remove(self.priorities);
t_remove(self.ids);
self.index[result_sync] = nil;
self.items[result_sync] = nil;
if size > k then
k = _percolate_up(self.priorities, k, self.ids, self.index);
_percolate_down(self.priorities, k, self.ids, self.index);
end
return result, item, result_sync;
end
function indexed_heap:remove(id)
return self:remove_index(self.index[id]);
end
local mt = { __index = indexed_heap };
local _M = {
create = function()
return setmetatable({
ids = {}; -- heap of ids, sync'd with priorities
items = {}; -- map id->items
priorities = {}; -- heap of priorities
index = {}; -- map of id->index of id in ids
current_id = 1.5
}, mt);
end
};
return _M;
|
