Add simple k-d-tree

master
Lars Mueller 2021-02-07 00:40:25 +01:00
parent 386f49a6ac
commit 155ab6deb7
2 changed files with 55 additions and 0 deletions

View File

@ -78,6 +78,7 @@ for _, component in ipairs{
"quaternion",
"minetest",
"trie",
"kdtree",
"heap",
"ranked_set",
"b3d"

54
kdtree.lua Normal file
View File

@ -0,0 +1,54 @@
local metatable = {__index = getfenv(1)}
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 self
end
function get_nearest_neighbor(self, vector)
local min_distance = math.huge
local nearest_neighbor
local distance_func = self.distance
local axis = tree.axis
local function visit(tree)
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] - color[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