1package main
 2
 3import "gonum.org/v1/gonum/spatial/kdtree"
 4
 5type galaxy []Location
 6
 7func (g galaxy) Len() int                              { return len(g) }
 8func (g galaxy) Index(i int) kdtree.Comparable         { return g[i] }
 9func (g galaxy) Slice(start, end int) kdtree.Interface { return g[start:end] }
10func (g galaxy) Pivot(d kdtree.Dim) int {
11	return plane{galaxy: g, Dim: d}.Pivot()
12}
13
14type plane struct {
15	kdtree.Dim
16	galaxy
17}
18
19func (p plane) Less(i, j int) bool {
20	switch p.Dim {
21	case 0:
22		return p.galaxy[i].a < p.galaxy[j].a
23	case 1:
24		return p.galaxy[i].e < p.galaxy[j].e
25	case 2:
26		return p.galaxy[i].m < p.galaxy[j].m
27	case 3:
28		return p.galaxy[i].v < p.galaxy[j].v
29	default:
30		panic("illegal dimension")
31	}
32}
33
34func (p plane) Pivot() int {
35	return kdtree.Partition(p, kdtree.MedianOfMedians(p))
36}
37
38func (p plane) Slice(start, end int) kdtree.SortSlicer {
39	p.galaxy = p.galaxy[start:end]
40	return p
41}
42
43func (p plane) Swap(i, j int) {
44	p.galaxy[i], p.galaxy[j] = p.galaxy[j], p.galaxy[i]
45}