64 lines
1.9 KiB
Lua
64 lines
1.9 KiB
Lua
-- Localize globals
|
|
local assert, math, modlib, setmetatable, table, unpack = assert, math, modlib, setmetatable, table, unpack
|
|
|
|
-- Set environment
|
|
local _ENV = {}
|
|
setfenv(1, _ENV)
|
|
|
|
local metatable = {__index = _ENV}
|
|
|
|
distance = modlib.vector.distance
|
|
|
|
--: vectors first vector is used to infer the dimension
|
|
--: distance (vector, other_vector) -> number, default: modlib.vector.distance
|
|
function new(vectors, distance)
|
|
assert(#vectors > 0, "vector list must not be empty")
|
|
local dimension = #vectors[1]
|
|
local function builder(vectors, axis)
|
|
if #vectors == 1 then return { value = vectors[1] } end
|
|
table.sort(vectors, function(a, b) return a[axis] > b[axis] end)
|
|
local median = math.floor(#vectors / 2)
|
|
local next_axis = ((axis + 1) % dimension) + 1
|
|
return setmetatable({
|
|
axis = axis,
|
|
pivot = vectors[median],
|
|
left = builder({ unpack(vectors, 1, median) }, next_axis),
|
|
right = builder({ unpack(vectors, median + 1) }, next_axis)
|
|
}, metatable)
|
|
end
|
|
local self = builder(vectors, 1)
|
|
self.distance = distance
|
|
return setmetatable(self, metatable)
|
|
end
|
|
|
|
function get_nearest_neighbor(self, vector)
|
|
local min_distance = math.huge
|
|
local nearest_neighbor
|
|
local distance_func = self.distance
|
|
local function visit(tree)
|
|
local axis = tree.axis
|
|
if tree.value ~= nil then
|
|
local distance = distance_func(tree.value, vector)
|
|
if distance < min_distance then
|
|
min_distance = distance
|
|
nearest_neighbor = tree.value
|
|
end
|
|
return
|
|
else
|
|
local this_side, other_side = tree.left, tree.right
|
|
if vector[axis] < tree.pivot[axis] then this_side, other_side = other_side, this_side end
|
|
visit(this_side)
|
|
if tree.pivot then
|
|
local dist = math.abs(tree.pivot[axis] - vector[axis])
|
|
if dist <= min_distance then visit(other_side) end
|
|
end
|
|
end
|
|
end
|
|
visit(self)
|
|
return nearest_neighbor, min_distance
|
|
end
|
|
|
|
-- TODO insertion & deletion + rebalancing
|
|
|
|
-- Export environment
|
|
return _ENV |