

 This page uses content from Wikipedia. The original article was at K-d tree. The list of authors can be seen in the page history. As with Rosetta Code, the text of Wikipedia is available under the GNU FDL. (See links for details on variance)

A k-d tree (short for k-dimensional tree) is a space-partitioning data structure for organizing points in a k-dimensional space. k-d trees are a useful data structure for several applications, such as searches involving a multidimensional search key (e.g. range searches and nearest neighbor searches).
k-d trees are a special case of binary space partitioning trees.

k-d trees are not suitable, however, for efficiently finding the nearest neighbor in high dimensional spaces. As a general rule, if the dimensionality is k, the number of points in the data, N, should be N ≫ 2k.
Otherwise, when k-d trees are used with high-dimensional data, most of the points in the tree will be evaluated and the efficiency is no better than exhaustive search, and other methods such as approximate nearest-neighbor are used instead.

Task: Construct a k-d tree and perform a nearest neighbor search for two example data sets:

 The Wikipedia example data of [(2,3), (5,4), (9,6), (4,7), (8,1), (7,2)].
 1000 3-d points uniformly distributed in a 3-d cube.
For the Wikipedia example, find the nearest neighbor to point (9, 2)
For the random data, pick a random location and find the nearest neighbor.

In addition, instrument your code to count the number of nodes visited in the nearest neighbor search. Count a node as visited if any field of it is accessed.

Output should show the point searched for, the point found,
the distance to the point, and the number of nodes visited.

There are variant algorithms for constructing the tree.
You can use a simple median strategy or implement something more efficient.
Variants of the nearest neighbor search include nearest N neighbors, approximate nearest neighbor, and range searches.
You do not have to implement these.
The requirement for this task is specifically the nearest single neighbor.
Also there are algorithms for inserting, deleting, and balancing k-d trees.
These are also not required for the task.

-- 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
