Task: K-d tree

C: #include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <time.h>

#define MAX_DIM 3
struct kd_node_t{
    double x[MAX_DIM];
    struct kd_node_t *left, *right;
};

    inline double
dist(struct kd_node_t *a, struct kd_node_t *b, int dim)
{
    double t, d = 0;
    while (dim--) {
        t = a->x[dim] - b->x[dim];
        d += t * t;
    }
    return d;
}
inline void swap(struct kd_node_t *x, struct kd_node_t *y) {
    double tmp[MAX_DIM];
    memcpy(tmp,  x->x, sizeof(tmp));
    memcpy(x->x, y->x, sizeof(tmp));
    memcpy(y->x, tmp,  sizeof(tmp));
}



    struct kd_node_t*
find_median(struct kd_node_t *start, struct kd_node_t *end, int idx)
{
    if (end <= start) return NULL;
    if (end == start + 1)
        return start;

    struct kd_node_t *p, *store, *md = start + (end - start) / 2;
    double pivot;
    while (1) {
        pivot = md->x[idx];

        swap(md, end - 1);
        for (store = p = start; p < end; p++) {
            if (p->x[idx] < pivot) {
                if (p != store)
                    swap(p, store);
                store++;
            }
        }
        swap(store, end - 1);

        
        if (store->x[idx] == md->x[idx])
            return md;

        if (store > md) end = store;
        else        start = store;
    }
}

    struct kd_node_t*
make_tree(struct kd_node_t *t, int len, int i, int dim)
{
    struct kd_node_t *n;

    if (!len) return 0;

    if ((n = find_median(t, t + len, i))) {
        i = (i + 1) % dim;
        n->left  = make_tree(t, n - t, i, dim);
        n->right = make_tree(n + 1, t + len - (n + 1), i, dim);
    }
    return n;
}


int visited;

void nearest(struct kd_node_t *root, struct kd_node_t *nd, int i, int dim,
        struct kd_node_t **best, double *best_dist)
{
    double d, dx, dx2;

    if (!root) return;
    d = dist(root, nd, dim);
    dx = root->x[i] - nd->x[i];
    dx2 = dx * dx;

    visited ++;

    if (!*best || d < *best_dist) {
        *best_dist = d;
        *best = root;
    }

    
    if (!*best_dist) return;

    if (++i >= dim) i = 0;

    nearest(dx > 0 ? root->left : root->right, nd, i, dim, best, best_dist);
    if (dx2 >= *best_dist) return;
    nearest(dx > 0 ? root->right : root->left, nd, i, dim, best, best_dist);
}

#define N 1000000
#define rand1() (rand() / (double)RAND_MAX)
#define rand_pt(v) { v.x[0] = rand1(); v.x[1] = rand1(); v.x[2] = rand1(); }
int main(void)
{
    int i;
    struct kd_node_t wp[] = {
        {{2, 3}}, {{5, 4}}, {{9, 6}}, {{4, 7}}, {{8, 1}}, {{7, 2}}
    };
    struct kd_node_t testNode = {{9, 2}};
    struct kd_node_t *root, *found, *million;
    double best_dist;

    root = make_tree(wp, sizeof(wp) / sizeof(wp[1]), 0, 2);

    visited = 0;
    found = 0;
    nearest(root, &testNode, 0, 2, &found, &best_dist);

    printf(">> WP tree\nsearching for (%g, %g)\n"
            "found (%g, %g) dist %g\nseen %d nodes\n\n",
            testNode.x[0], testNode.x[1],
            found->x[0], found->x[1], sqrt(best_dist), visited);

    million =(struct kd_node_t*) calloc(N, sizeof(struct kd_node_t));
    srand(time(0));
    for (i = 0; i < N; i++) rand_pt(million[i]);

    root = make_tree(million, N, 0, 3);
    rand_pt(testNode);

    visited = 0;
    found = 0;
    nearest(root, &testNode, 0, 3, &found, &best_dist);

    printf(">> Million tree\nsearching for (%g, %g, %g)\n"
            "found (%g, %g, %g) dist %g\nseen %d nodes\n",
            testNode.x[0], testNode.x[1], testNode.x[2],
            found->x[0], found->x[1], found->x[2],
            sqrt(best_dist), visited);

    
    int sum = 0, test_runs = 100000;
    for (i = 0; i < test_runs; i++) {
        found = 0;
        visited = 0;
        rand_pt(testNode);
        nearest(root, &testNode, 0, 3, &found, &best_dist);
        sum += visited;
    }
    printf("\n>> Million tree\n"
            "visited %d nodes for %d random findings (%f per lookup)\n",
            sum, test_runs, sum/(double)test_runs);

    

    return 0;
}




Go: 


package main

import (
    "fmt"
    "math"
    "math/rand"
    "sort"
    "time"
)


type point []float64


func (p point) sqd(q point) float64 {
    var sum float64
    for dim, pCoord := range p {
        d := pCoord - q[dim]
        sum += d * d
    }
    return sum
}




type kdNode struct {
    domElt      point
    split       int
    left, right *kdNode
}   

type kdTree struct {
    n      *kdNode
    bounds hyperRect
}
    
type hyperRect struct {
    min, max point
}



func (hr hyperRect) copy() hyperRect {
    return hyperRect{append(point{}, hr.min...), append(point{}, hr.max...)}
}   
    



func newKd(pts []point, bounds hyperRect) kdTree {
    var nk2 func([]point, int) *kdNode
    nk2 = func(exset []point, split int) *kdNode {
        if len(exset) == 0 {
            return nil
        }
        
        
        
        sort.Sort(part{exset, split})
        m := len(exset) / 2
        d := exset[m]
        for m+1 < len(exset) && exset[m+1][split] == d[split] {
            m++
        }
        
        s2 := split + 1
        if s2 == len(d) {
            s2 = 0
        }
        return &kdNode{d, split, nk2(exset[:m], s2), nk2(exset[m+1:], s2)}
    }
    return kdTree{nk2(pts, 0), bounds}
}



type part struct {
    pts   []point
    dPart int
}


func (p part) Len() int { return len(p.pts) }
func (p part) Less(i, j int) bool {
    return p.pts[i][p.dPart] < p.pts[j][p.dPart]
}
func (p part) Swap(i, j int) { p.pts[i], p.pts[j] = p.pts[j], p.pts[i] }





func (t kdTree) nearest(p point) (best point, bestSqd float64, nv int) {
    return nn(t.n, p, t.bounds, math.Inf(1))
}



func nn(kd *kdNode, target point, hr hyperRect,
    maxDistSqd float64) (nearest point, distSqd float64, nodesVisited int) {
    if kd == nil {
        return nil, math.Inf(1), 0
    }
    nodesVisited++
    s := kd.split
    pivot := kd.domElt
    leftHr := hr.copy()
    rightHr := hr.copy()
    leftHr.max[s] = pivot[s]
    rightHr.min[s] = pivot[s]
    targetInLeft := target[s] <= pivot[s]
    var nearerKd, furtherKd *kdNode
    var nearerHr, furtherHr hyperRect
    if targetInLeft {
        nearerKd, nearerHr = kd.left, leftHr
        furtherKd, furtherHr = kd.right, rightHr
    } else {
        nearerKd, nearerHr = kd.right, rightHr
        furtherKd, furtherHr = kd.left, leftHr
    }
    var nv int
    nearest, distSqd, nv = nn(nearerKd, target, nearerHr, maxDistSqd)
    nodesVisited += nv
    if distSqd < maxDistSqd {
        maxDistSqd = distSqd
    }
    d := pivot[s] - target[s]
    d *= d
    if d > maxDistSqd {
        return
    }
    if d = pivot.sqd(target); d < distSqd {
        nearest = pivot
        distSqd = d
        maxDistSqd = distSqd
    }
    tempNearest, tempSqd, nv := nn(furtherKd, target, furtherHr, maxDistSqd)
    nodesVisited += nv
    if tempSqd < distSqd {
        nearest = tempNearest
        distSqd = tempSqd
    }
    return
}

func main() {
    rand.Seed(time.Now().Unix())
    kd := newKd([]point{{2, 3}, {5, 4}, {9, 6}, {4, 7}, {8, 1}, {7, 2}},
        hyperRect{point{0, 0}, point{10, 10}})
    showNearest("WP example data", kd, point{9, 2})
    kd = newKd(randomPts(3, 1000), hyperRect{point{0, 0, 0}, point{1, 1, 1}})
    showNearest("1000 random 3d points", kd, randomPt(3))
}   
    
func randomPt(dim int) point {
    p := make(point, dim)
    for d := range p {
        p[d] = rand.Float64()
    }
    return p
}   
    
func randomPts(dim, n int) []point {
    p := make([]point, n)
    for i := range p {
        p[i] = randomPt(dim) 
    } 
    return p
}
    
func showNearest(heading string, kd kdTree, p point) {
    fmt.Println()
    fmt.Println(heading)
    fmt.Println("point:           ", p)
    nn, ssq, nv := kd.nearest(p)
    fmt.Println("nearest neighbor:", nn)
    fmt.Println("distance:        ", math.Sqrt(ssq))
    fmt.Println("nodes visited:   ", nv)
}