Ebo/BucketKDTree

The tree I use in Tahoe.

BucketKdTree?.java

package de._4geeks.robots.guns.DC;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;

public class BucketKdTree<T extends Info> {

	private int dim;
	private int bucketSize;
	private int entries;
	
	private Bucket buckets;

	public BucketKdTree(int dim, int bucketSize) {
		this.dim = dim;
		this.bucketSize = bucketSize;

		this.buckets = new Bucket(bucketSize, 0);
		this.entries = 0;
	}

	public void addInfo(T info) {
		buckets.addInfo(this, info);
		entries += 1;
	}

	public void clear() {
		this.buckets = new Bucket(bucketSize, 0);
		this.entries = 0;
	}

	public void balance() {
		ArrayList<T> l = collect();
		this.buckets = new Bucket(bucketSize, 0);
		this.buckets.balance(l, this);
	}
	
	public ArrayList<T> collect() {
		ArrayList<T> l = new ArrayList<T>();
		buckets.collect(l);
		return l;
	}
	
	public ArrayList<T> getBucket(double[] features) {
		return buckets.getBucket(features);
	}

	class Bucket {
		ArrayList<T> b;
		int dim;

		Bucket top;
		Bucket bottom;
		double border;

		public Bucket(int bucketSize, int dim) {
			this.dim = dim;
			b = new ArrayList<T>(bucketSize);
		}

		public ArrayList<T> getBucket(double[] features) {
			if (b != null) {
				return b;
			} else {
				if (features[dim] >= border) {
					return top.getBucket(features);
				} else {
					return bottom.getBucket(features);
				}
			}
		}

		public void addInfo(BucketKdTree<T> tree, T info) {
			if (b != null) {
				if (b.size() >= tree.bucketSize) {
					subdivide(tree);
					if (info.getFeature(dim) >= border) {
						top.addInfo(tree, info);
					} else {
						bottom.addInfo(tree, info);
					}
				} else {
					b.add(info);
				}
			} else {
				if (info.getFeature(dim) >= border) {
					top.addInfo(tree, info);
				} else {
					bottom.addInfo(tree, info);
				}
			}

		}

		private void subdivide(BucketKdTree<T> tree) {

			// This calculates the median for the bucket.

			Collections.sort(b, new InfoComparator(dim));
			border = b.get(b.size() / 2).getFeature(dim);
			
			//System.out.println(String.format("Subdivide: %d %f", dim, border));
			
			top = new Bucket(bucketSize, (dim+1) % tree.getDim());
			bottom = new Bucket(bucketSize, (dim+1) % tree.getDim());
			for (T i : b) {
				if (i.getFeature(dim) >= border) {
					top.addInfo(tree, i);
				} else {
					bottom.addInfo(tree, i);
				}
			}
			b = null;

		}
		
		public void collect(ArrayList<T> l) {
			if(b != null) {
				l.addAll(b);
			} else {
				top.collect(l);
				bottom.collect(l);
			}
		}
		
		public void balance(ArrayList<T> l, BucketKdTree<T> tree) {
			if(l.size() >= tree.bucketSize) {
				Collections.sort(l, new InfoComparator(dim));
				border = l.get(l.size() / 2).getFeature(dim);
				
				top = new Bucket(bucketSize, (dim+1) % tree.getDim());
				bottom = new Bucket(bucketSize, (dim+1) % tree.getDim());	
				ArrayList<T> topList = new ArrayList<T>();
				ArrayList<T> bottomList = new ArrayList<T>();
				
				for(T i : l) {
					if (i.getFeature(dim) >= border) {
						topList.add(i);
					} else {
						bottomList.add(i);
					}
				}
				
				top.balance(topList, tree);
				bottom.balance(bottomList, tree);
				b = null;
			} else {
				b = l;
			}
				
		}
		
		public int size() {
			return b.size();
		}

		class InfoComparator implements Comparator<T> {
			int i;

			public InfoComparator(int i) {
				this.i = i;
			}

			public int compare(T o1, T o2) {
				return Double.compare(o1.getFeature(i), o2.getFeature(i));
			}
		}

	}

	public int getDim() {
		return dim;
	}

	public int getBucketSize() {
		return bucketSize;
	}

	public int getEntries() {
		return entries;
	}
}

Info.java

package de._4geeks.robots.guns.DC;

public interface Info {
	public double getFeature(int i);
}

[BSD License]