(phixonline)-->
-- demo\rosetta\kd_tree.exw
with javascript_semantics
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)
return sum(sq_power(sq_sub(p,q),2))
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)
while true do
atom 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 */
atom sx = kd_nodes[stor][X][idx],
mx = kd_nodes[md][X][idx]
if sx == mx then exit end if
if stor>md then last = stor else first = stor end if
end while
return md
end function
function make_tree(object t, len, i, dim)
if sequence(t) then
kd_nodes = deep_copy(t)
t = 1
end if
integer n = iff(len=0?0: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]
if dx<=0 then {l,r} = {r,l} end if
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
// aside: shave a little off under pwa/p2js, but 1e6 build still takes ~20s
constant test_runs = iff(platform()=JS?10_000:100_000)
procedure test(string id, sequence nodes, 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)
printf(1,"\n>> %s tree\nsearching for %v\nfound %v dist %g\nseen %d nodes\n",
{id, test_node[X], kd_nodes[best][X], 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
--
integer total = 0
for i=1 to test_runs do
best = 0
visited = 0
test_node = rand_pt()
nearest(root, test_node, 1, 3)
total += visited
end for
printf(1,"average behaviour: "&
"visited %d nodes for %,d random findings (%f per lookup)\n",
{total, test_runs, total/test_runs});
end if
end procedure
atom t0 = time()
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)
?elapsed(time()-t0)
if platform()!=JS then
?"done"
{} = wait_key()
end if