153 lines
4.3 KiB
Text
153 lines
4.3 KiB
Text
enum X,LEFT,RIGHT -- (keeping the IDX on each node too would not be a bad idea..)
|
|
-- (the code below deduces it from the (unbalanced) tree depth)
|
|
|
|
sequence kd_nodes -- nb probably not best coding style
|
|
|
|
function sqdist(sequence p,q)
|
|
atom dsq = sum(sq_power(sq_sub(p,q),2))
|
|
return dsq
|
|
end function
|
|
|
|
procedure swap(integer x,y)
|
|
{kd_nodes[x],kd_nodes[y]} = {kd_nodes[y],kd_nodes[x]}
|
|
end procedure
|
|
|
|
function find_median(integer first, last, idx)
|
|
if last<=first then return NULL end if
|
|
if last==first+1 then return first end if
|
|
|
|
integer p, stor, md = first+floor((last-first)/2)
|
|
atom pivot
|
|
while true do
|
|
pivot = kd_nodes[md][X][idx]
|
|
|
|
swap(md, last-1)
|
|
stor = first
|
|
for p=first to last-1 do
|
|
if kd_nodes[p][X][idx]<pivot then
|
|
if p!=stor then
|
|
swap(p, stor)
|
|
end if
|
|
stor += 1
|
|
end if
|
|
end for
|
|
swap(stor, last-1)
|
|
|
|
/* median has duplicate values */
|
|
if kd_nodes[stor][X][idx] == kd_nodes[md][X][idx] then
|
|
return md
|
|
end if
|
|
if stor>md then last = stor else first = stor end if
|
|
end while
|
|
end function
|
|
|
|
function make_tree(object t, len, i, dim)
|
|
if sequence(t) then kd_nodes = t; t = 1; end if
|
|
if len=0 then return 0 end if
|
|
integer n = find_median(t, t+len, i)
|
|
if n then
|
|
i = mod(i,dim)+1
|
|
if length(kd_nodes[n])!=1 then ?9/0 end if -- (once)
|
|
kd_nodes[n] = {kd_nodes[n][X],0,0} -- (add l/r slots)
|
|
kd_nodes[n][LEFT] = make_tree(t, n-t, i, dim)
|
|
kd_nodes[n][RIGHT] = make_tree(n+1, t+len-(n+1), i, dim)
|
|
end if
|
|
return n
|
|
end function
|
|
|
|
integer visited,
|
|
best
|
|
atom best_dist
|
|
|
|
procedure nearest(integer root, sequence nd, integer i, dim)
|
|
if root=0 then return end if
|
|
atom d = sqdist(kd_nodes[root][X],nd[X]),
|
|
dx = kd_nodes[root][X][i] - nd[X][i],
|
|
dx2 = dx * dx;
|
|
|
|
visited += 1
|
|
|
|
if best=0 or d<best_dist then
|
|
best_dist = d;
|
|
best = root;
|
|
end if
|
|
|
|
/* if chance of exact match is high */
|
|
if best_dist=0 then return end if
|
|
|
|
i = mod(i,dim)+1
|
|
integer {l,r} = kd_nodes[root][LEFT..RIGHT]
|
|
{l,r} = iff(dx>0?{l,r}:{r,l})
|
|
nearest(l, nd, i, dim)
|
|
if (dx2 >= best_dist) then return end if
|
|
nearest(r, nd, i, dim)
|
|
end procedure
|
|
|
|
function rand_pt() return sq_rnd({{0,0,0}}) end function
|
|
|
|
function make_3D()
|
|
sequence s = {}
|
|
for i=0 to 9 do
|
|
for j=0 to 9 do
|
|
for k=0 to 9 do
|
|
s = append(s,{{i,j,k}})
|
|
end for
|
|
end for
|
|
end for
|
|
return s
|
|
end function
|
|
|
|
constant test_runs = 100_000
|
|
|
|
procedure test(string id, sequence nodes, sequence test_node, bool show_average_behaviour=false)
|
|
integer dim = length(nodes[1][X]),
|
|
root = make_tree(nodes, length(nodes), 1, dim)
|
|
best = 0
|
|
visited = 0
|
|
nearest(root, test_node, 1, dim)
|
|
string tn = sprint(test_node[X]),
|
|
bn = sprint(kd_nodes[best][X])
|
|
printf(1,"\n>> %s tree\nsearching for %s\n"&
|
|
"found %s dist %g\nseen %d nodes\n",
|
|
{id, tn, bn, sqrt(best_dist), visited})
|
|
|
|
if show_average_behaviour then
|
|
--
|
|
-- search many random points to see average behavior.
|
|
--
|
|
-- tree size vs avg nodes visited:
|
|
-- 10 ~ 7
|
|
-- 100 ~ 16.5
|
|
-- 1000 ~ 25.5
|
|
-- 10000 ~ 32.8
|
|
-- 100000 ~ 38.3
|
|
-- 1000000 ~ 42.6
|
|
-- 10000000 ~ 46.7
|
|
--
|
|
int tot = 0
|
|
for i=1 to test_runs do
|
|
best = 0
|
|
visited = 0
|
|
test_node = rand_pt()
|
|
nearest(root, test_node, 1, 3)
|
|
tot += visited
|
|
end for
|
|
printf(1,"average behaviour: "&
|
|
"visited %d nodes for %,d random findings (%f per lookup)\n",
|
|
{tot, test_runs, tot/test_runs});
|
|
end if
|
|
end procedure
|
|
|
|
sequence wp = {{{2, 3}}, {{5, 4}}, {{9, 6}}, {{4, 7}}, {{8, 1}}, {{7, 2}}},
|
|
test_node = {{9, 2}}
|
|
test("WP",wp,test_node)
|
|
|
|
sequence cube = make_3D()
|
|
test_node = sq_mul(rand_pt(),10)
|
|
test("3D",cube,test_node,true)
|
|
|
|
constant N = 1_000_000
|
|
sequence million = repeat(0,N)
|
|
for i=1 to N do million[i] = rand_pt() end for
|
|
test_node = rand_pt()
|
|
test("Million",million,test_node,true)
|