modlib/ranked_set.lua
2021-06-17 19:45:08 +02:00

321 lines
7.4 KiB
Lua

-- Localize globals
local assert, modlib, pairs, setmetatable, table = assert, modlib, pairs, setmetatable, table
-- Set environment
local _ENV = {}
setfenv(1, _ENV)
local metatable = {__index = _ENV}
comparator = modlib.table.default_comparator
--+ Uses a weight-balanced binary tree
function new(comparator)
return setmetatable({comparator = comparator, root = {total = 0}}, metatable)
end
function len(self)
return self.root.total
end
metatable.__len = len
function is_empty(self)
return len(self) == 0
end
local function insert_all(tree, _table)
if tree.left then
insert_all(tree.left, _table)
end
table.insert(_table, tree.key)
if tree.right then
insert_all(tree.right, _table)
end
end
function to_table(self)
local table = {}
if not is_empty(self) then
insert_all(self.root, table)
end
return table
end
--> iterator: function() -> `rank, key` with ascending rank
function ipairs(self, min, max)
if is_empty(self) then
return function() end
end
min = min or 1
local tree = self.root
local current_rank = (tree.left and tree.left.total or 0) + 1
repeat
if min == current_rank then
break
end
local left, right = tree.left, tree.right
if min < current_rank then
current_rank = current_rank - (left and left.right and left.right.total or 0) - 1
tree = left
else
current_rank = current_rank + (right and right.left and right.left.total or 0) + 1
tree = right
end
until not tree
max = max or len(self)
local to_visit = {tree}
tree = nil
local rank = min - 1
local function next()
if not tree then
local len = #to_visit
if len == 0 then return end
tree = to_visit[len]
to_visit[len] = nil
else
while tree.left do
table.insert(to_visit, tree)
tree = tree.left
end
end
local key = tree.key
tree = tree.right
return key
end
return function()
if rank >= max then
return
end
local key = next()
if key == nil then
return
end
rank = rank + 1
return rank, key
end
end
local function _right_rotation(parent, right, left)
local new_parent = parent[left]
parent[left] = new_parent[right]
new_parent[right] = parent
parent.total = (parent[left] and parent[left].total or 0) + (parent[right] and parent[right].total or 0) + 1
assert(parent.total > 0 or (parent.left == nil and parent.right == nil))
new_parent.total = (new_parent[left] and new_parent[left].total or 0) + parent.total + 1
return new_parent
end
local function right_rotation(parent)
return _right_rotation(parent, "right", "left")
end
local function left_rotation(parent)
return _right_rotation(parent, "left", "right")
end
local function _rebalance(parent)
local left_count, right_count = (parent.left and parent.left.total or 0), (parent.right and parent.right.total or 0)
if right_count > 1 and left_count * 2 < right_count then
return left_rotation(parent)
end
if left_count > 1 and right_count * 2 < left_count then
return right_rotation(parent)
end
return parent
end
-- Rebalances a parent chain
local function rebalance(self, len, parents, sides)
if len <= 1 then
return
end
for i = len, 2, -1 do
parents[i] = _rebalance(parents[i])
parents[i - 1][sides[i - 1]] = parents[i]
end
self.root = parents[1]
end
local function _insert(self, key, replace)
assert(key ~= nil)
if is_empty(self) then
self.root = {key = key, total = 1}
return
end
local comparator = self.comparator
local parents, sides = {}, {}
local tree = self.root
repeat
local tree_key = tree.key
local compared = comparator(key, tree_key)
if compared == 0 then
if replace then
tree.key = key
return tree_key
end
return
end
table.insert(parents, tree)
local side = compared < 0 and "left" or "right"
table.insert(sides, side)
tree = tree[side]
until not tree
local len = #parents
parents[len][sides[len]] = {key = key, total = 1}
for _, parent in pairs(parents) do
parent.total = parent.total + 1
end
rebalance(self, len, parents, sides)
end
function insert(self, key)
return _insert(self, key)
end
function insert_or_replace(self, key)
return _insert(self, key, true)
end
local function _delete(self, key, is_rank)
assert(key ~= nil)
if is_empty(self) then
return
end
local comparator = self.comparator
local parents, sides = {}, {}
local tree = self.root
local rank = (tree.left and tree.left.total or 0) + 1
repeat
local tree_key = tree.key
local compared
if is_rank then
if key == rank then
compared = 0
elseif key < rank then
rank = rank - (tree.left and tree.left.right and tree.left.right.total or 0) - 1
compared = -1
else
rank = rank + (tree.right and tree.right.left and tree.right.left.total or 0) + 1
compared = 1
end
else
compared = comparator(key, tree_key)
end
if compared == 0 then
local len = #parents
local left, right = tree.left, tree.right
if left then
tree.total = tree.total - 1
if right then
-- Obtain successor
local side = left.total > right.total and "left" or "right"
local other_side = side == "left" and "right" or "left"
local sidemost = tree[side]
while sidemost[other_side] do
sidemost.total = sidemost.total - 1
table.insert(parents, sidemost)
table.insert(sides, other_side)
sidemost = sidemost[other_side]
end
-- Replace deleted key
tree.key = sidemost.key
-- Replace the successor by it's single child
parents[len][sides[len]] = sidemost[side]
else
if len == 0 then
self.root = left or {total = 0}
else
parents[len][sides[len]] = left
end
end
elseif right then
if len == 0 then
self.root = right or {total = 0}
else
tree.total = tree.total - 1
parents[len][sides[len]] = right
end
else
if len == 0 then
self.root = {total = 0}
else
parents[len][sides[len]] = nil
end
end
for _, parent in pairs(parents) do
parent.total = parent.total - 1
end
rebalance(self, len, parents, sides)
if is_rank then
return tree_key
end
return rank, tree_key
end
table.insert(parents, tree)
local side
if compared < 0 then
side = "left"
else
side = "right"
end
table.insert(sides, side)
tree = tree[side]
until not tree
end
function delete(self, key)
return _delete(self, key)
end
delete_by_key = delete
function delete_by_rank(self, rank)
return _delete(self, rank, true)
end
--> `rank, key` if the key was found
--> `rank` the key would have if inserted
function get(self, key)
if is_empty(self) then return end
local comparator = self.comparator
local tree = self.root
local rank = (tree.left and tree.left.total or 0) + 1
while tree do
local compared = comparator(key, tree.key)
if compared == 0 then
return rank, tree.key
end
if compared < 0 then
rank = rank - (tree.left and tree.left.right and tree.left.right.total or 0) - 1
tree = tree.left
else
rank = rank + (tree.right and tree.right.left and tree.right.left.total or 0) + 1
tree = tree.right
end
end
return rank
end
get_by_key = get
--> key
function get_by_rank(self, rank)
local tree = self.root
local current_rank = (tree.left and tree.left.total or 0) + 1
repeat
if rank == current_rank then
return tree.key
end
local left, right = tree.left, tree.right
if rank < current_rank then
current_rank = current_rank - (left and left.right and left.right.total or 0) - 1
tree = left
else
current_rank = current_rank + (right and right.left and right.left.total or 0) + 1
tree = right
end
until not tree
end
-- Export environment
return _ENV