Binary Search Tree

binary-search-tree.lua

function BinarySearchTree(cmp)
    local root
    local size = 0
    local function has(v, el)
        if not v then
            return false
        end
        if v.value == el then
            return true
        end
        if cmp(el, v.value) then
            return has(v.left, el)
        else
            return has(v.right, el)
        end
    end
    local function add(v, el)
        if not v then
            size = size + 1
            return { value = el }
        end
        if v.value ~= el then
            if cmp(el, v.value) then
                v.left = add(v.left, el)
            else
                v.right = add(v.right, el)
            end
        end
        return v
    end
    local function pop(v)
        if not v then
            return nil, nil
        end
        local value
        if v.left then
            v.left, value = pop(v.left)
            return v, value
        elseif v.right then
            return v.right, v.value
        else
            return nil, v.value
        end
    end
    local function del(v, el)
        if not v then
            return nil
        end
        if v.value == el then
            size = size - 1
            if v.left and v.right then
                v.right, v.value = pop(v.right)
                return v
            elseif v.left then
                return v.left
            elseif v.right then
                return v.right
            else
                return nil
            end
        else
            if cmp(el, v.value) then
                v.left = del(v.left, el)
            else
                v.right = del(v.right, el)
            end
            return v
        end
    end
    local function peek(v)
        if v == nil then
            return nil
        end
        if v.left then
            return peek(v.left)
        else
            return v.value
        end
    end
    local function dump(v, s)
        if (v == nil) then
            return
        end
        print(s .. tostring(v.value))
        dump(v.left, s .. ' ')
        dump(v.right, s .. ' ')
    end
    local function it(v)
        if not v then
            return
        end
        it(v.left)
        coroutine.yield(v.value)
        it(v.right)
    end
    return {
        size = function () return size end,
        has = function (el) return has(root, el) end,
        add = function (el) root = add(root, el) end,
        del = function (el) root = del(root, el) end,
        pop = function ()
            if size == 0 then
                return nil
            end
            local value
            size = size - 1
            root, value = pop(root)
            return value
        end,
        peek = function () return peek(root) end,
        dump = function (s) return dump(root, s or '') end,
        it = function ()
            local co = coroutine.create(function () it(root) end)
            return function()
                local _, value = coroutine.resume(co)
                return value
            end
        end
    }
end