-- Red-black trees: -- Every node is red or black -- Root is black -- All leaves (nils) are black -- If a node is red, both its children are black -- For each node, all paths to descendant leaves contain the same number of black nodes IntervalTree = {} function IntervalTree.create() local RED = true local BLACK = false local sentinel = {} sentinel.parent = sentinel sentinel.left = sentinel sentinel.right = sentinel sentinel.color = BLACK local tree = { root = sentinel } local function successor(node) if node.right ~= sentinel then node = node.right while node.left ~= sentinel do node = node.left end return node end local y = node.parent while y ~= sentinel and node == y.right do node = y y = y.parent end return y end local function findNode(start, stop) local node = tree.root while node ~= sentinel do if start < node.start then node = node.left elseif start > node.start then node = node.right elseif stop < node.stop then node = node.left elseif stop > node.stop then node = node.right else return node end end return sentinel end local function find(node, num) if node == sentinel or num > node.max then return {} end local results = {} for _, result in ipairs(find(node.left, num)) do table.insert(results, result) end if num >= node.start then if num <= node.stop then table.insert(results, node.value) end for _, result in ipairs(find(node.right, num)) do table.insert(results, result) end end return results end function tree.find(num) return find(tree.root, num) end function tree.findInterval(start, stop) return findNode(start, stop).value end local function rotateLeft(node) local x = node local y = x.right x.right = y.left if y.left ~= sentinel then y.left.parent = x end y.parent = x.parent if x.parent == sentinel then tree.root = y elseif x == x.parent.left then x.parent.left = y else x.parent.right = y end y.left = x x.parent = y x.max = math.max(x.stop, x.left.max or 0, x.right.max or 0) y.max = math.max(y.stop, y.left.max or 0, x.right.max or 0) end local function rotateRight(node) local x = node local y = x.left x.left = y.right if y.right ~= sentinel then y.right.parent = x end y.parent = x.parent if x.parent == sentinel then tree.root = y elseif x == x.parent.right then x.parent.right = y else x.parent.left = y end y.right = x x.parent = y x.max = math.max(x.stop, x.left.max or 0, x.right.max or 0) y.max = math.max(y.stop, y.left.max or 0, x.right.max or 0) end local function insertFixup(node) local y while node.parent.color == RED do if node.parent == node.parent.parent.left then y = node.parent.parent.right if y.color == RED then node.parent.color = BLACK y.color = BLACK node.parent.parent.color = RED node = node.parent.parent elseif node == node.parent.right then node = node.parent rotateLeft(node) else node.parent.color = BLACK node.parent.parent.color = RED rotateRight(node.parent.parent) end else y = node.parent.parent.left if y.color == RED then node.parent.color = BLACK y.color = BLACK node.parent.parent.color = RED node = node.parent.parent elseif node == node.parent.left then node = node.parent rotateRight(node) else node.parent.color = BLACK node.parent.parent.color = RED rotateLeft(node.parent.parent) end end end tree.root.color = BLACK end function tree.insert(start, stop, value) local y = sentinel local x = tree.root while x ~= sentinel do y = x if start < x.start then x = x.left elseif start > x.start then x = x.right elseif stop < x.stop then x = x.left elseif stop > x.stop then x = x.right else for key, value in pairs(value) do if not x.value[key] then x.value[key] = {} end for k, v in pairs(value) do x.value[key][k] = v end end return end end local node = { start = start, stop = stop, max = stop, value = value, parent = y, left = sentinel, right = sentinel, color = RED, } if y == sentinel then tree.root = node elseif start < y.start then y.left = node else y.right = node end insertFixup(node) while node.parent ~= sentinel do if node.max > node.parent.max then node.parent.max = node.max end node = node.parent end end return tree end