kdtree fixes

This commit is contained in:
Lars Mueller 2021-02-07 14:00:41 +01:00
parent 155ab6deb7
commit 006beff2ef

@ -21,15 +21,15 @@ function new(vectors, distance)
end end
local self = builder(vectors, 1) local self = builder(vectors, 1)
self.distance = distance self.distance = distance
return self return setmetatable(self, metatable)
end end
function get_nearest_neighbor(self, vector) function get_nearest_neighbor(self, vector)
local min_distance = math.huge local min_distance = math.huge
local nearest_neighbor local nearest_neighbor
local distance_func = self.distance local distance_func = self.distance
local axis = tree.axis
local function visit(tree) local function visit(tree)
local axis = tree.axis
if tree.value ~= nil then if tree.value ~= nil then
local distance = distance_func(tree.value, vector) local distance = distance_func(tree.value, vector)
if distance < min_distance then if distance < min_distance then
@ -42,7 +42,7 @@ function get_nearest_neighbor(self, vector)
if vector[axis] < tree.pivot[axis] then this_side, other_side = other_side, this_side end if vector[axis] < tree.pivot[axis] then this_side, other_side = other_side, this_side end
visit(this_side) visit(this_side)
if tree.pivot then if tree.pivot then
local dist = math.abs(tree.pivot[axis] - color[axis]) local dist = math.abs(tree.pivot[axis] - vector[axis])
if dist <= min_distance then visit(other_side) end if dist <= min_distance then visit(other_side) end
end end
end end