2 local setmetatable
= setmetatable
;
3 local math_floor
= math
.floor;
4 local t_remove
= table.remove;
6 local function _heap_insert(self
, item
, sync
, item2
, index
)
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
;
21 local function _percolate_up(self
, k
, sync
, index
)
23 local tmp_sync
= sync
[k
];
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
];
38 local function _percolate_down(self
, k
, sync
, index
)
40 local tmp_sync
= sync
[k
];
44 if child
~= size
and self
[child
] > self
[child
+ 1] then
47 if tmp
> self
[child
] then
48 self
[k
] = self
[child
];
49 sync
[k
] = sync
[child
];
64 local function _heap_pop(self
, sync
, index
)
66 if size
== 0 then return nil; end
68 local result
= self
[1];
69 local result_sync
= sync
[1];
70 index
[result_sync
] = nil;
74 return result
, result_sync
;
76 self
[1] = t_remove(self
);
77 sync
[1] = t_remove(sync
);
80 _percolate_down(self
, 1, sync
, index
);
82 return result
, result_sync
;
85 local indexed_heap
= {};
87 function indexed_heap
:insert(item
, priority
, id
)
90 self
.current_id
= id
+ 1;
92 self
.items
[id
] = item
;
93 _heap_insert(self
.priorities
, priority
, self
.ids
, id
, self
.index
);
96 function indexed_heap
:pop()
97 local priority
, id
= _heap_pop(self
.priorities
, self
.ids
, self
.index
);
99 local item
= self
.items
[id
];
100 self
.items
[id
] = nil;
101 return priority
, item
, id
;
104 function indexed_heap
:peek()
105 return self
.priorities
[1];
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
;
112 k
= _percolate_up(self
.priorities
, k
, self
.ids
, self
.index
);
113 _percolate_down(self
.priorities
, k
, self
.ids
, self
.index
);
115 function indexed_heap
:remove_index(k
)
116 local result
= self
.priorities
[k
];
117 if result
== nil then return; end
119 local result_sync
= self
.ids
[k
];
120 local item
= self
.items
[result_sync
];
121 local size
= #self
.priorities
;
123 self
.priorities
[k
] = self
.priorities
[size
];
124 self
.ids
[k
] = self
.ids
[size
];
125 self
.index
[self
.ids
[k]]
= k
;
127 t_remove(self
.priorities
);
130 self
.index
[result_sync
] = nil;
131 self
.items
[result_sync
] = nil;
134 k
= _percolate_up(self
.priorities
, k
, self
.ids
, self
.index
);
135 _percolate_down(self
.priorities
, k
, self
.ids
, self
.index
);
138 return result
, item
, result_sync
;
140 function indexed_heap
:remove(id
)
141 return self
:remove_index(self
.index
[id
]);
144 local mt
= { __index
= indexed_heap
};
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