kdTree nearest

package s3;


import java.awt.Color;

/*************************************************************************
 *************************************************************************/

import edu.princeton.cs.algs4.In;
import edu.princeton.cs.algs4.Out;
import edu.princeton.cs.algs4.Point2D;
import edu.princeton.cs.algs4.RectHV;
import edu.princeton.cs.algs4.StdDraw;

public class KdTree {
    private Node root;

    private static class Node {
        private Point2D p; // the point
        private RectHV rect; // the axis-aligned rectangle corresponding to this node
        private Node lb; // the left/bottom subtree
        private Node rt; // the right/top subtree

        public Node(Point2D point){
            this.p = point;
        }
    }
    // construct an empty set of points
    public KdTree() {
    }

    // is the set empty?
    public boolean isEmpty() {
        return root == null;
    }

    // number of points in the set
    public int size() {
        return size(root);
    }

    private int size(Node node) {
        if (node == null)
            return 0;
        else
            return 1 + size(node.lb) + size(node.rt);
    }

    // add the point p to the set (if it is not already in the set)
    public void insert(Point2D p) {
        root = insertAt(root, p, true);
    }

    private Node insertAt(Node node, Point2D p, boolean isX) {
        if (node == null)
            return new Node(p);

        double dist = Double.MAX_VALUE;

        if (isX)
            dist = node.p.x() - p.x();
        else
            dist = node.p.y() - p.y();

        if (dist > 0)
            node.lb = insertAt(node.lb, p, !isX);
        else if (dist < 0)
            node.rt = insertAt(node.rt, p, !isX);
        else if (node.p.equals(p))
            node.p = p;
        else{
            dist = node.p.compareTo(p);
            if (dist > 0)
                insertAt(node.lb, p, !isX);
            else
                insertAt(node.rt, p, !isX);
        }

        return node;
    }

    // does the set contain the point p?
    public boolean contains(Point2D p) {
        return containsAt(root, p);
    }

    private boolean containsAt(Node node, Point2D p) {
        if (node == null)
            return false;
        if (node.p.equals(p))
            return true;
        else if (p.compareTo(node.p) < 0)
            return containsAt(node.lb, p);
        else
            return containsAt(node.rt, p);
    }

    // draw all of the points to standard draw
    public void draw() {
        if (isEmpty())
            return;
        else
            drawlb(root, 0, new Point2D(0, 0));
    }

    private void drawlb(Node node, int lvl, Point2D last) {
        if (node == null)
            return;

        if (lvl % 2 == 0){
            StdDraw.setPenColor(Color.red);
            StdDraw.filledCircle(node.p.x(), node.p.y(), 0.005);
            StdDraw.line(node.p.x(), last.y(), node.p.x(), 0);
        }
        else{
            StdDraw.setPenColor(Color.blue);
            StdDraw.filledCircle(node.p.x(), node.p.y(), 0.005);
            StdDraw.line(0, node.p.y(), last.x(), node.p.y());
        }

        drawlb(node.lb, lvl + 1, node.p);
        drawrt(node.rt, lvl + 1, node.p);
    }

    private void drawrt(Node node, int lvl, Point2D last) {
        if (node == null)
            return;

        if (lvl % 2 == 0){
            StdDraw.setPenColor(Color.red);
            StdDraw.filledCircle(node.p.x(), node.p.y(), 0.005);
            StdDraw.line(node.p.x(), 1, node.p.x(), last.y());
        }
        else{
            StdDraw.setPenColor(Color.blue);
            StdDraw.filledCircle(node.p.x(), node.p.y(), 0.005);
            StdDraw.line(last.x(), node.p.y(), 1, node.p.y());
        }

        drawlb(node.lb, lvl + 1, node.p);
        drawrt(node.rt, lvl + 1, node.p);
    }

    // all points in the set that are inside the rectangle
    public Iterable<Point2D> range(RectHV rect) {
        return null;
    }

    // a nearest neighbor in the set to p; null if set is empty
    public Point2D nearest(Point2D p) {
        if (isEmpty())
            return null;
        return nearest(root, p, true, root.p);
    }

    private Point2D nearest(Node node, Point2D p, boolean isX, Point2D npoint) {
        double dist = Double.MAX_VALUE;
        if (node.p.distanceTo(p) < npoint.distanceTo(p))
            npoint = node.p;

        if(isX)
            dist = node.p.x() - p.x();
        else
            dist = node.p.y() - p.y();


        if (node.p.equals(p))
            return node.p;

        if ((node.lb != null) && (node.rt != null)){
            if (dist > 0){
                    return nearest(node.lb, p, !isX, npoint);
            }
            else{
                    return nearest(node.rt, p, !isX, npoint);
            }
        }
        else if ((node.lb != null) && (node.rt == null)){
            if (dist > 0)
                return nearest(node.lb, p, !isX, npoint);
            else
                return npoint;
        }
        else if ((node.lb == null) && (node.rt != null)){
            if (dist > 0)
                return npoint;
            else
                return nearest(node.rt, p, !isX, npoint);
        }
        else
            return npoint;
    }

    public static void main(String[] args) {
        In in = new In();
        Out out = new Out();
        int N = in.readInt(), C = in.readInt(), T = 50;
        Point2D[] queries = new Point2D[C];
        KdTree tree = new KdTree();
        out.printf("Inserting %d points into tree\n", N);
        for (int i = 0; i < N; i++) {
            tree.insert(new Point2D(in.readDouble(), in.readDouble()));
        }
        out.printf("tree.size(): %d\n", tree.size());
        out.printf("Testing `nearest` method, querying %d points\n", C);

        for (int i = 0; i < C; i++) {
            queries[i] = new Point2D(in.readDouble(), in.readDouble());
            out.printf("%s: %s\n", queries[i], tree.nearest(queries[i]));
        }
        for (int i = 0; i < T; i++) {
            for (int j = 0; j < C; j++) {
                tree.nearest(queries[j]);
            }
        }
    }

    /*
    public static void main(String[] args) {
        In in = new In();
        Out out = new Out();
        int N = in.readInt(), C = in.readInt(), T = 20;
        KdTree tree = new KdTree();
        Point2D [] points = new Point2D[C];
        out.printf("Inserting %d points into tree\n", N);
        for (int i = 0; i < N; i++) {
            tree.insert(new Point2D(in.readDouble(), in.readDouble()));
        }
        out.printf("tree.size(): %d\n", tree.size());
        out.printf("Testing contains method, querying %d points\n", C);
        for (int i = 0; i < C; i++) {
            points[i] = new Point2D(in.readDouble(), in.readDouble());
            out.printf("%s: %s\n", points[i], tree.contains(points[i]));
        }
        for (int i = 0; i < T; i++) {
            for (int j = 0; j < C; j++) {
                tree.contains(points[j]);
            }
        }
    }*/
}