Это полная реализация KD-Tree, я использовал несколько библиотек для хранения точек и прямоугольников. Эти библиотеки находятся в свободном доступе. С этими классами можно сделать мои собственные классы для хранения точки и прямоугольника. Пожалуйста, поделитесь своим отзывом.
import java.util.ArrayList;
import java.util.List;
import edu.princeton.cs.algs4.In;
import edu.princeton.cs.algs4.Point2D;
import edu.princeton.cs.algs4.RectHV;
import edu.princeton.cs.algs4.StdDraw;
public class KdTree {
private static class Node {
public Point2D point; // the point
public RectHV rect; // the axis-aligned rectangle corresponding to this
public Node lb; // the left/bottom subtree
public Node rt; // the right/top subtree
public int size;
public double x = 0;
public double y = 0;
public Node(Point2D p, RectHV rect, Node lb, Node rt) {
super();
this.point = p;
this.rect = rect;
this.lb = lb;
this.rt = rt;
x = p.x();
y = p.y();
}
}
private Node root = null;;
public KdTree() {
}
public boolean isEmpty() {
return root == null;
}
public int size() {
return rechnenSize(root);
}
private int rechnenSize(Node node) {
if (node == null) {
return 0;
} else {
return node.size;
}
}
public void insert(Point2D p) {
if (p == null) {
throw new NullPointerException();
}
if (isEmpty()) {
root = insertInternal(p, root, 0);
root.rect = new RectHV(0, 0, 1, 1);
} else {
root = insertInternal(p, root, 1);
}
}
// at odd level we will compare x coordinate, and at even level we will
// compare y coordinate
private Node insertInternal(Point2D pointToInsert, Node node, int level) {
if (node == null) {
Node newNode = new Node(pointToInsert, null, null, null);
newNode.size = 1;
return newNode;
}
if (level % 2 == 0) {//Horizontal partition line
if (pointToInsert.y() < node.y) {//Traverse in bottom area of partition
node.lb = insertInternal(pointToInsert, node.lb, level + 1);
if(node.lb.rect == null){
node.lb.rect = new RectHV(node.rect.xmin(), node.rect.ymin(),
node.rect.xmax(), node.y);
}
} else {//Traverse in top area of partition
if (!node.point.equals(pointToInsert)) {
node.rt = insertInternal(pointToInsert, node.rt, level + 1);
if(node.rt.rect == null){
node.rt.rect = new RectHV(node.rect.xmin(), node.y,
node.rect.xmax(), node.rect.ymax());
}
}
}
} else if (level % 2 != 0) {//Vertical partition line
if (pointToInsert.x() < node.x) {//Traverse in left area of partition
node.lb = insertInternal(pointToInsert, node.lb, level + 1);
if(node.lb.rect == null){
node.lb.rect = new RectHV(node.rect.xmin(), node.rect.ymin(),
node.x, node.rect.ymax());
}
} else {//Traverse in right area of partition
if (!node.point.equals(pointToInsert)) {
node.rt = insertInternal(pointToInsert, node.rt, level + 1);
if(node.rt.rect == null){
node.rt.rect = new RectHV(node.x, node.rect.ymin(),
node.rect.xmax(), node.rect.ymax());
}
}
}
}
node.size = 1 + rechnenSize(node.lb) + rechnenSize(node.rt);
return node;
}
public boolean contains(Point2D p) {
return containsInternal(p, root, 1);
}
private boolean containsInternal(Point2D pointToSearch, Node node, int level) {
if (node == null) {
return false;
}
if (level % 2 == 0) {//Horizontal partition line
if (pointToSearch.y() < node.y) {
return containsInternal(pointToSearch, node.lb, level + 1);
} else {
if (node.point.equals(pointToSearch)) {
return true;
}
return containsInternal(pointToSearch, node.rt, level + 1);
}
} else {//Vertical partition line
if (pointToSearch.x() < node.x) {
return containsInternal(pointToSearch, node.lb, level + 1);
} else {
if (node.point.equals(pointToSearch)) {
return true;
}
return containsInternal(pointToSearch, node.rt, level + 1);
}
}
}
public void draw() {
StdDraw.clear();
drawInternal(root, 1);
}
private void drawInternal(Node node, int level) {
if (node == null) {
return;
}
StdDraw.setPenColor(StdDraw.BLACK);
StdDraw.setPenRadius(0.02);
node.point.draw();
double sx = node.rect.xmin();
double ex = node.rect.xmax();
double sy = node.rect.ymin();
double ey = node.rect.ymax();
StdDraw.setPenRadius(0.01);
if (level % 2 == 0) {
StdDraw.setPenColor(StdDraw.BLUE);
sy = ey = node.y;
} else {
StdDraw.setPenColor(StdDraw.RED);
sx = ex = node.x;
}
StdDraw.line(sx, sy, ex, ey);
drawInternal(node.lb, level + 1);
drawInternal(node.rt, level + 1);
}
/**
* Find the points which lies in the rectangle as parameter
* @param rect
* @return
*/
public Iterable<Point2D> range(RectHV rect) {
List<Point2D> resultList = new ArrayList<Point2D>();
rangeInternal(root, rect, resultList);
return resultList;
}
private void rangeInternal(Node node, RectHV rect, List<Point2D> resultList) {
if (node == null) {
return;
}
if (node.rect.intersects(rect)) {
if (rect.contains(node.point)) {
resultList.add(node.point);
}
rangeInternal(node.lb, rect, resultList);
rangeInternal(node.rt, rect, resultList);
}
}
public Point2D nearest(Point2D p) {
if(root == null){
return null;
}
Champion champion = new Champion(root.point,Double.MAX_VALUE);
return nearestInternal(p, root, champion, 1).champion;
}
private Champion nearestInternal(Point2D targetPoint, Node node,
Champion champion, int level) {
if (node == null) {
return champion;
}
double dist = targetPoint.distanceSquaredTo(node.point);
int newLevel = level + 1;
if (dist < champion.championDist) {
champion.champion = node.point;
champion.championDist = dist;
}
boolean goLeftOrBottom = false;
//We will decide which part to be visited first, based upon in which part point lies.
//If point is towards left or bottom part, we traverse in that area first, and later on decide
//if we need to search in other part too.
if(level % 2 == 0){
if(targetPoint.y() < node.y){
goLeftOrBottom = true;
}
} else {
if(targetPoint.x() < node.x){
goLeftOrBottom = true;
}
}
if(goLeftOrBottom){
nearestInternal(targetPoint, node.lb, champion, newLevel);
Point2D orientationPoint = createOrientationPoint(node.x,node.y,targetPoint,level);
double orientationDist = orientationPoint.distanceSquaredTo(targetPoint);
//We will search on the other part only, if the point is very near to partitioned line
//and champion point found so far is far away from the partitioned line.
if(orientationDist < champion.championDist){
nearestInternal(targetPoint, node.rt, champion, newLevel);
}
} else {
nearestInternal(targetPoint, node.rt, champion, newLevel);
Point2D orientationPoint = createOrientationPoint(node.x,node.y,targetPoint,level);
//We will search on the other part only, if the point is very near to partitioned line
//and champion point found so far is far away from the partitioned line.
double orientationDist = orientationPoint.distanceSquaredTo(targetPoint);
if(orientationDist < champion.championDist){
nearestInternal(targetPoint, node.lb, champion, newLevel);
}
}
return champion;
}
/**
* Returns the point from a partitioned line, which can be directly used to calculate
* distance between partitioned line and the target point for which neighbours are to be searched.
* @param linePointX
* @param linePointY
* @param targetPoint
* @param level
* @return
*/
private Point2D createOrientationPoint(double linePointX, double linePointY, Point2D targetPoint, int level){
if(level % 2 == 0){
return new Point2D(targetPoint.x(),linePointY);
} else {
return new Point2D(linePointX,targetPoint.y());
}
}
private static class Champion{
public Point2D champion;
public double championDist;
public Champion(Point2D c, double d){
champion = c;
championDist = d;
}
}
public static void main(String[] args) {
String filename = "/home/raman/Downloads/kdtree/circle100.txt";
In in = new In(filename);
KdTree kdTree = new KdTree();
while (!in.isEmpty()) {
double x = in.readDouble();
double y = in.readDouble();
Point2D p = new Point2D(x, y);
kdTree.insert(p);
}
// kdTree.print();
System.out.println(kdTree.size());
kdTree.draw();
System.out.println(kdTree.nearest(new Point2D(0.4, 0.5)));
System.out.println(new Point2D(0.7, 0.4).distanceSquaredTo(new Point2D(0.9,0.5)));
System.out.println(new Point2D(0.7, 0.4).distanceSquaredTo(new Point2D(0.9,0.4)));
}
}