aboutsummaryrefslogtreecommitdiffstats
path: root/util/indexedbheap.lua
blob: 7f193d5409a1a8077f826b666cd367c5c3e85c25 (plain)
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;