A very basic B-Tree implementation in lua

btree.lua

-- This is a simple implementation of BTree Currently including only the hard part: insertions and deletions
--
-- This implementation supports using custom comparison functions (this allows to include more data in the keys for example).
-- I also designed this implementation to allow easy support for on-disk storage.
--
-- Copyright © 2025 ascpial
-- This project is licensed under the [MIT License](https://mit-license.org/).


--- @class BTreeNode
--- @field tree BTree
--- @field children BTreeNode[]
--- @field elements table
--- @field root boolean
local BTreeNode = {}
local BTreeNode_mt = { __index = BTreeNode }

--- @return BTreeNode
function BTreeNode.new(tree, elements, children, root)
  local node = {
    tree = tree,
    elements = elements or {},
    children = children or {},
    root = root or false,
  }
  setmetatable(node, BTreeNode_mt)
  return node
end

function BTreeNode:isLeaf()
  return #self.children == 0
end

function BTreeNode:insert(elt, i)
  local nl, nr
  i = i or 1
  if not self:isLeaf() then
    if self.elements[i] then
      if self.tree.comp(elt, self.elements[i]) then
        elt, nl, nr = self.tree:getNode(self.children[i]):insert(elt)
      else
        return self:insert(elt, i + 1)
      end
    else
      elt, nl, nr = self.tree:getNode(self.children[i]):insert(elt)
    end
  end
  if elt then
    if self.elements[i] and self.tree.comp(self.elements[i], elt) then
      return self:insert(elt, i + 1)
    end
    table.insert(self.elements, i, elt)
    table.insert(self.children, i, nl)
    self.children[i + 1] = nr -- nil if the node is a leaf
    if #self.elements > self.tree.m then
      -- We need to split the array in two
      local n = #self.elements
      local mn = math.ceil(n / 2)
      local nodeLeft = self.tree:createNode({ table.unpack(self.elements, 1, mn - 1) },
        { table.unpack(self.children, 1, mn) })
      local nodeRight = self.tree:createNode({ table.unpack(self.elements, mn + 1, n) },
        { table.unpack(self.children, mn + 1, n + 1) })
      local median = self.elements[mn]
      if self.root then
        self.root = false
        return self.tree:getNode(self.tree:createNode({ median }, { nodeLeft, nodeRight }, true))
      else
        return median, nodeLeft, nodeRight
      end
    end
  end
  if self.root then
    return self
  end
end

function BTreeNode:min()
  if self:isLeaf() then
    return self.elements[0]
  else
    return self.tree:getNode(self.children)[0]:min()
  end
end

function BTreeNode:max()
  if self:isLeaf() then
    return self.elements[#self.elements]
  else
    return self.tree:getNode(self.children[#self.children]):max()
  end
end

function BTreeNode:minElts()
  if self.root then
    return 0
  elseif self:isLeaf() then
    return math.ceil(self.tree.m / 2)
  else
    return math.floor(self.tree.m / 2)
  end
end

--- @return BTreeNode | nil
function BTreeNode:delete(elt, i)
  i = i or 1
  local child
  if self.elements[i] then
    if self.tree.comp(self.elements[i], elt) and not self.tree.comp(elt, self.elements[i]) then
      return self:delete(elt, i + 1)
    elseif self.tree.comp(elt, self.elements[i]) and self.tree.comp(self.elements[i], elt) then -- elements are equal
      if self:isLeaf() then
        table.remove(self.elements, i)
        return self
      else
        local max = self.tree:getNode(self.children[i]):max()
        self.elements[i] = max
        elt = max
      end
    end
  else
    if self:isLeaf() then
      return nil
    end
  end
  child = self.tree:getNode(self.children[i])
  if child:delete(elt) == nil then -- did not find the element
    return
  end
  if #child.elements < child:minElts() then
    local leftChild = self.tree:getNode(self.children[i - 1])
    local rightChild = self.tree:getNode(self.children[i + 1])
    if leftChild and #leftChild.elements > leftChild:minElts() then
      table.insert(child.elements, 1, self.elements[i - 1])
      table.insert(child.children, 1, (table.remove(leftChild.children)))
      self.elements[i - 1] = table.remove(leftChild.elements)
    elseif rightChild and #rightChild.elements > rightChild:minElts() then
      table.insert(child.elements, self.elements[i])
      table.insert(child.children, (table.remove(rightChild.children, 1)))
      self.elements[i] = table.remove(rightChild.elements, 1)
    else
      local j = (leftChild and i) or (i + 1)
      local left = self.tree:getNode(self.children[j - 1])
      local right = self.tree:getNode(self.children[j])
      table.insert(left.elements, table.remove(self.elements, j - 1))
      for _, v in ipairs(right.elements) do
        table.insert(left.elements, v)
      end
      for _, v in ipairs(right.children) do
        table.insert(left.children, v)
      end
      table.remove(self.children, j)
    end
  end
  if self.root and #self.elements == 0 then
    local newRoot = self.tree:getNode(self.children[1])
    newRoot.root = true
    self.root = false
    return newRoot
  end
  return self
end

--- @class BTree
--- @field m number The max number of elements per node
--- @field root BTreeNode
--- @field comp function
--- @field equal function
local BTree = {}
local BTree_mt = { __index = BTree }

--- @param m number The maximum number of elements in a node
--- @param comp function The comparison function, e.g. <=. Make sure that comp(a, b) and comp(b, a) is equivalent to a == b
--- @return BTree
function BTree.new(m, comp)
  local o = {
    m = m,
    comp = comp,
  }
  o.root = BTreeNode.new(o, {}, {}, true)
  setmetatable(o, BTree_mt)
  return o
end

--- @param id any An ID corresponding to a node, the node by default
--- @return BTreeNode The BTreeNode associated with id
function BTree:getNode(id)
  return id
end

--- @param elements any[] The seperators keys
--- @param children any[] The IDs of the children of the new node
--- @param root boolean? Wether this element is the root or not
--- @return any The ID of the newly created node
function BTree:createNode(elements, children, root)
  return BTreeNode.new(self, elements, children, root)
end

--- @return any The minimum value according to the comparison function
function BTree:min() return self.root:min() end

--- @return any The maximum value according to the comparison function
function BTree:max() return self.root:max() end

--- @param elt any The element to insert in the tree
function BTree:insert(elt)
  self.root = self.root:insert(elt)
end

--- @param elt any The element to remove from the tree
--- @return boolean Wether the element was present in the tree or not
function BTree:delete(elt)
  local newRoot = self.root:delete(elt)
  if newRoot then
    self.root = newRoot
    return true
  else
    return false
  end
end

return BTree