Skip to content

Instantly share code, notes, and snippets.

@themattchan

themattchan/KdTree.java

Last active Mar 3, 2016
Embed
What would you like to do?
Coursera Algs pt 1 Hw 5, example for my blog post.
import edu.princeton.cs.algs4.*;
import java.util.*;
public class KdTree {
private abstract class Node implements TreeApply {
protected Point2D point;
public abstract void apply (TreeFunction function);
}
private interface TreeFunction {
public void match(VertNode node);
public void match(HorizNode node);
}
private interface TreeApply {
public void apply(TreeFunction function);
}
private final class VertNode extends Node {
private HorizNode lt, rt;
public VertNode (Point2D p) {
point = p;
N++;
}
@Override
public void apply(TreeFunction function) {
function.match(this);
}
}
private final class HorizNode extends Node {
private VertNode lt, rt;
public HorizNode (Point2D p) {
point = p;
N++;
}
@Override
public void apply(TreeFunction function) {
function.match(this);
}
}
private int N; // number of points in tree
private VertNode root;
public KdTree() {
N = 0;
root = null;
}
public boolean isEmpty() {
return root == null && N == 0;
}
public int size() {
return N;
}
private class InsertPoint implements TreeFunction {
private Point2D p;
public InsertPoint(Point2D p) {
this.p = p;
}
@Override
public void match(VertNode node) {
if (node == null) return;
if (p.equals(node.point)) return;
else if (p.x() < node.point.x()) {
if (node.lt == null)
node.lt = new HorizNode(p);
else match(node.lt);
}
else {
if (node.rt == null)
node.rt = new HorizNode(p);
else match(node.rt);
}
}
@Override
public void match(HorizNode node) {
if (node == null) return;
if (p.equals(node.point)) return;
else if (p.x() < node.point.y()) {
if (node.lt == null)
node.lt = new VertNode(p);
else match(node.lt);
}
else {
if (node.rt == null)
node.rt = new VertNode(p);
else match(node.rt);
}
}
}
public void insert(Point2D p) {
if (root == null)
root = new VertNode(p);
else
root.apply(new InsertPoint(p));
}
private class ContainsPoint implements TreeFunction {
private Point2D p;
private boolean result;
public ContainsPoint(Point2D p) {
this.p = p;
}
public boolean getResult() {
return result;
}
@Override
public void match(VertNode node) {
if (node == null) {
result = false;
}
else if (p.equals(node.point)) {
result = true;
}
else if (p.x() < node.point.x()) {
if (node.lt == null)
result = false;
else match(node.lt);
}
else {
if (node.rt == null)
result = false;
else match(node.rt);
}
}
@Override
public void match(HorizNode node) {
if (node == null) {
result = false;
}
if (p.equals(node.point)) {
result = true;
}
else if (p.x() < node.point.y()) {
if (node.lt == null)
result = false;
else match(node.lt);
}
else {
if (node.rt == null)
result = false;
else match(node.rt);
}
}
}
public boolean contains(Point2D p) {
if (isEmpty()) return false;
ContainsPoint function = new ContainsPoint(p);
root.apply(function);
return function.getResult();
}
private class DrawKdTree implements TreeFunction {
private static final double PEN_RADIUS = 0.016d;
@Override
public void match(VertNode node) {
if (node == null) return;
match(node.lt);
match(node.rt);
if (node.point == null) return;
StdDraw.setPenColor(StdDraw.BLACK);
StdDraw.setPenRadius(PEN_RADIUS);
StdDraw.point(node.point.x(), node.point.y());
StdDraw.setPenRadius();
StdDraw.setPenColor(StdDraw.RED);
StdDraw.line(node.point.x(), 0,
node.point.x(), 1);
}
@Override
public void match(HorizNode node) {
if (node == null) return;
match(node.lt);
match(node.rt);
if (node.point == null) return;
StdDraw.setPenColor(StdDraw.BLACK);
StdDraw.setPenRadius(PEN_RADIUS);
StdDraw.point(node.point.x(), node.point.y());
StdDraw.setPenRadius();
StdDraw.setPenColor(StdDraw.BLUE);
StdDraw.line(0, node.point.y(),
1, node.point.y());
}
}
public void draw() {
if (isEmpty()) return;
root.apply(new DrawKdTree());
}
private class FindRange implements TreeFunction {
private RectHV rect;
private List<Point2D> points = new ArrayList<>();
public FindRange(RectHV rect) {
this.rect = rect;
}
public List<Point2D> getResult() { return points; }
@Override
public void match(VertNode n) {
if (n == null) {
return;
}
else if (n.point.x() < rect.xmin()) {
match(n.rt);
}
else if (n.point.x() > rect.xmax()) {
match(n.lt);
}
else {
points.add(n.point);
match(n.rt);
match(n.lt);
}
}
@Override
public void match(HorizNode n) {
if (n == null) {
return;
}
else if (n.point.y() < rect.ymin()) {
match(n.rt);
}
else if (n.point.y() > rect.ymax()) {
match(n.lt);
}
else {
points.add(n.point);
match(n.rt);
match(n.lt);
}
}
}
public Iterable<Point2D> range(RectHV rect) {
if (isEmpty()) return null;
FindRange function = new FindRange(rect);
root.apply(function);
return function.getResult();
}
private class FindNearestPoint implements TreeFunction {
private Point2D p, nearest;
public FindNearestPoint(Point2D p) {
this.p = p;
this.nearest == null;
}
private Double getDistance() {
if (nearest == null) return Double.POSITIVE_INFINITY;
return p.distanceSquaredTo(nearest);
}
public Point2D getResult() { return nearest; }
@Override
public void match(VertNode n) {
if (n == null) return;
if (
}
@Override
public void match(HorizNode n) {
if (n == null) return;
}
}
public Point2D nearest(Point2D p) {
if (isEmpty()) return null;
FindNearestPoint function = new FindNearestPoint(p);
root.apply(function);
return function.getResult();
}
public static void main(String[] args) {}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment