util/indexedbheap.lua

   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 result = self.priorities[k];
 117         if result == nil then return; end
 118
 119         local result_sync = self.ids[k];
 120         local item = self.items[result_sync];
 121         local size = #self.priorities;
 122
 123         self.priorities[k] = self.priorities[size];
 124         self.ids[k] = self.ids[size];
 125         self.index[self.ids[k]] = k;
 126
 127         t_remove(self.priorities);
 128         t_remove(self.ids);
 129
 130         self.index[result_sync] = nil;
 131         self.items[result_sync] = nil;
 132
 133         if size > k then
 134                 k = _percolate_up(self.priorities, k, self.ids, self.index);
 135                 k = _percolate_down(self.priorities, k, self.ids, self.index);
 136         end
 137
 138         return result, item, result_sync;
 139 end
 140 function indexed_heap:remove(id)
 141         return self:remove_index(self.index[id]);
 142 end
 143
 144 local mt = { __index = indexed_heap };
 145
 146 local _M = {
 147         create = function()
 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
 153                         current_id = 1.5
 154                 }, mt);
 155         end
 156 };
 157 return _M;