Task: K-d tree




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




Python: from random import seed, random
from time import time
from operator import itemgetter
from collections import namedtuple
from math import sqrt
from copy import deepcopy


def sqd(p1, p2):
    return sum((c1 - c2) ** 2 for c1, c2 in zip(p1, p2))


class KdNode(object):
    __slots__ = ("dom_elt", "split", "left", "right")

    def __init__(self, dom_elt, split, left, right):
        self.dom_elt = dom_elt
        self.split = split
        self.left = left
        self.right = right


class Orthotope(object):
    __slots__ = ("min", "max")

    def __init__(self, mi, ma):
        self.min, self.max = mi, ma


class KdTree(object):
    __slots__ = ("n", "bounds")

    def __init__(self, pts, bounds):
        def nk2(split, exset):
            if not exset:
                return None
            exset.sort(key=itemgetter(split))
            m = len(exset) // 2
            d = exset[m]
            while m + 1 < len(exset) and exset[m + 1][split] == d[split]:
                m += 1
            d = exset[m]


            s2 = (split + 1) % len(d)  
            return KdNode(d, split, nk2(s2, exset[:m]),
                                    nk2(s2, exset[m + 1:]))
        self.n = nk2(0, pts)
        self.bounds = bounds

T3 = namedtuple("T3", "nearest dist_sqd nodes_visited")


def find_nearest(k, t, p):
    def nn(kd, target, hr, max_dist_sqd):
        if kd is None:
            return T3([0.0] * k, float("inf"), 0)

        nodes_visited = 1
        s = kd.split
        pivot = kd.dom_elt
        left_hr = deepcopy(hr)
        right_hr = deepcopy(hr)
        left_hr.max[s] = pivot[s]
        right_hr.min[s] = pivot[s]

        if target[s] <= pivot[s]:
            nearer_kd, nearer_hr = kd.left, left_hr
            further_kd, further_hr = kd.right, right_hr
        else:
            nearer_kd, nearer_hr = kd.right, right_hr
            further_kd, further_hr = kd.left, left_hr

        n1 = nn(nearer_kd, target, nearer_hr, max_dist_sqd)
        nearest = n1.nearest
        dist_sqd = n1.dist_sqd
        nodes_visited += n1.nodes_visited

        if dist_sqd < max_dist_sqd:
            max_dist_sqd = dist_sqd
        d = (pivot[s] - target[s]) ** 2
        if d > max_dist_sqd:
            return T3(nearest, dist_sqd, nodes_visited)
        d = sqd(pivot, target)
        if d < dist_sqd:
            nearest = pivot
            dist_sqd = d
            max_dist_sqd = dist_sqd

        n2 = nn(further_kd, target, further_hr, max_dist_sqd)
        nodes_visited += n2.nodes_visited
        if n2.dist_sqd < dist_sqd:
            nearest = n2.nearest
            dist_sqd = n2.dist_sqd

        return T3(nearest, dist_sqd, nodes_visited)

    return nn(t.n, p, t.bounds, float("inf"))


def show_nearest(k, heading, kd, p):
    print(heading + ":")
    print("Point:           ", p)
    n = find_nearest(k, kd, p)
    print("Nearest neighbor:", n.nearest)
    print("Distance:        ", sqrt(n.dist_sqd))
    print("Nodes visited:   ", n.nodes_visited, "\n")


def random_point(k):
    return [random() for _ in range(k)]


def random_points(k, n):
    return [random_point(k) for _ in range(n)]

if __name__ == "__main__":
    seed(1)
    P = lambda *coords: list(coords)
    kd1 = KdTree([P(2, 3), P(5, 4), P(9, 6), P(4, 7), P(8, 1), P(7, 2)],
                  Orthotope(P(0, 0), P(10, 10)))
    show_nearest(2, "Wikipedia example data", kd1, P(9, 2))

    N = 400000
    t0 = time()
    kd2 = KdTree(random_points(3, N), Orthotope(P(0, 0, 0), P(1, 1, 1)))
    t1 = time()
    text = lambda *parts: "".join(map(str, parts))
    show_nearest(2, text("k-d tree with ", N,
                         " random 3D points (generation time: ",
                         t1-t0, "s)"),
                 kd2, random_point(3))