Algorithms-I : 5-kdtrees: implement KdTree

1 files changed, 202 insertions, 5 deletions

diff --git a/Algorithms/Part-I/5-KdTrees/KdTree.java b/Algorithms/Part-I/5-KdTrees/KdTree.java
index 877a2c3..73d9a90 100644
--- a/Algorithms/Part-I/5-KdTrees/KdTree.java
+++ b/Algorithms/Part-I/5-KdTrees/KdTree.java
@@ -2,49 +2,246 @@
 
 public class KdTree
 {
+    private static final double XMIN = 0;
+    private static final double YMIN = 0;
+    private static final double XMAX = 1;
+    private static final double YMAX = 1;
+
+    private Node root;
+    private int size;
+
+    private class Node {
+        private Point2D p;
+        private Node left, right;
+
+        public Node(Point2D p)
+        {
+            this.p = p;
+        }
+    }
+
+    private class NearestChampion {
+        private Point2D p;
+        private double d;
+        public NearestChampion(Point2D p, double d)
+        {
+            this.p = p;
+            this.d = d;
+        }
+    }
+
     // construct an empty tree of points
     public KdTree()
     {
+        size = 0;
+        root = null;
     }
 
     // is the tree empty?
     public boolean isEmpty()
     {
-        return true;
+        return (size == 0);
     }
 
     // number of points in the tree
     public int size()
     {
-        return 0;
+        return size;
     }
 
     // add the point p to the tree (if it is not already in the tree)
     public void insert(Point2D p)
     {
+        if (p == null) return;
+        if (root == null)
+        {
+            root = new Node(p);
+            size = 1;
+        }
+        else if (put(root, p, true))
+            size += 1;
+    }
+
+    private boolean put(Node n, Point2D p, boolean vsplit) {
+
+        int cmp;
+        if (vsplit)
+            cmp = Point2D.X_ORDER.compare(p, n.p);
+        else
+            cmp = Point2D.Y_ORDER.compare(p, n.p);
+
+        if (cmp < 0)
+        {
+            // add to left subtree
+            if (n.left != null)
+                return put(n.left, p, !vsplit);
+            n.left = new Node(p);
+            return true;
+        }
+
+        if (cmp == 0)
+        {
+            if (vsplit)
+            {
+                if (Point2D.Y_ORDER.compare(p, n.p) == 0)
+                    return false;
+            }
+            else if (Point2D.X_ORDER.compare(p, n.p) == 0)
+                return false;
+        }
+
+        // add to right subtree
+        if (n.right != null)
+            return put(n.right, p, !vsplit);
+        n.right = new Node(p);
+        return true;
     }
 
     // does the tree contain the point p?
     public boolean contains(Point2D p)
     {
-        return false;
+        if (root == null) return false;
+        return get(root, p, true);
+    }
+
+    private boolean get(Node n, Point2D p, boolean vsplit) {
+        if (p.compareTo(n.p) == 0) return true;
+        int cmp;
+        if (vsplit)
+            cmp = Point2D.X_ORDER.compare(p, n.p);
+        else
+            cmp = Point2D.Y_ORDER.compare(p, n.p);
+
+        if (cmp < 0)
+            return get(n.left, p, !vsplit);
+        return get(n.right, p, !vsplit);
     }
 
     // draw all of the points to standard draw
     public void draw()
     {
+        StdDraw.setPenRadius();
+        StdDraw.setPenColor(StdDraw.BLACK);
+        StdDraw.rectangle(XMAX/2.0, YMAX/2.0, XMAX/2.0, YMAX/2.0);
+        draw(root, true, new RectHV(XMIN, YMIN, XMAX, YMAX));
+    }
+
+    private void draw(Node n, boolean vsplit, RectHV r)
+    {
+        if (n == null) return;
+        if (vsplit)
+        {
+            StdDraw.setPenRadius();
+            StdDraw.setPenColor(StdDraw.RED);
+            StdDraw.line(n.p.x(), r.ymin(), n.p.x(), r.ymax());
+            if (n.left != null)
+                draw(n.left, !vsplit,
+                        new RectHV(r.xmin(), r.ymin(), n.p.x(), r.ymax()));
+            if (n.right != null)
+                draw(n.right, !vsplit,
+                        new RectHV(n.p.x(), r.ymin(), r.xmax(), r.ymax()));
+        }
+        else
+        {
+            StdDraw.setPenRadius();
+            StdDraw.setPenColor(StdDraw.BLUE);
+            StdDraw.line(r.xmin(), n.p.y(), r.xmax(), n.p.y());
+            if (n.left != null)
+                draw(n.left, !vsplit,
+                        new RectHV(r.xmin(), r.ymin(), r.xmax(), n.p.y()));
+            if (n.right != null)
+                draw(n.right, !vsplit,
+                        new RectHV(r.xmin(), n.p.y(), r.xmax(), r.ymax()));
+        }
+        StdDraw.setPenRadius(.01);
+        StdDraw.setPenColor(StdDraw.BLACK);
+        n.p.draw();
     }
 
     // all points in the tree that are inside the rectangle
     public Iterable<Point2D> range(RectHV rect)
     {
-        return null;
+        Stack<Point2D> stack = new Stack<Point2D>();
+
+        range(root, rect, stack, true);
+
+        return stack;
+    }
+
+    private void range(Node n, RectHV r, Stack<Point2D> s, boolean vsplit)
+    {
+        if (r.contains(n.p))
+            s.push(n.p);
+        if (vsplit)
+        {
+            if (n.left != null && r.xmin() < n.p.x())
+                range(n.left, r, s, !vsplit);
+            if (n.right != null && r.xmax() >= n.p.x())
+                range(n.right, r, s, !vsplit);
+        }
+        else
+        {
+            if (n.left != null && r.ymin() < n.p.y())
+                range(n.left, r, s, !vsplit);
+            if (n.right != null && r.ymax() >= n.p.y())
+                range(n.right, r, s, !vsplit);
+        }
     }
 
     // a nearest neighbor in the tree to p; null if tree is empty
     public Point2D nearest(Point2D p)
     {
-        return null;
+        if (p == null) return null;
+        if (root == null) return null;
+        NearestChampion ncp = new NearestChampion(null, Double.MAX_VALUE);
+        nearest(root, p, ncp, true);
+        return ncp.p;
+    }
+
+    private void nearest(Node n, Point2D p, NearestChampion ncp, boolean vsplit)
+    {
+        double d2 = p.distanceTo(n.p);
+        if (d2 < ncp.d)
+        {
+            ncp.d = d2;
+            ncp.p = n.p;
+        }
+        if (vsplit)
+        {
+            double nx = n.p.x();
+            if (p.x() < nx)
+            {
+                if (n.left != null)
+                    nearest(n.left, p, ncp, !vsplit);
+                if (n.right != null && ((nx - p.x()) < ncp.d))
+                    nearest(n.right, p, ncp, !vsplit);
+            }
+            else
+            {
+                if (n.right != null)
+                    nearest(n.right, p, ncp, !vsplit);
+                if (n.left != null && ((p.x() - nx) < ncp.d))
+                    nearest(n.left, p, ncp, !vsplit);
+            }
+        }
+        else
+        {
+            double ny = n.p.y();
+            if (p.y() < ny)
+            {
+                if (n.left != null)
+                    nearest(n.left, p, ncp, !vsplit);
+                if (n.right != null && ((ny - p.y()) < ncp.d))
+                    nearest(n.right, p, ncp, !vsplit);
+            }
+            else
+            {
+                if (n.right != null)
+                    nearest(n.right, p, ncp, !vsplit);
+                if (n.left != null && ((p.y() - ny) < ncp.d))
+                    nearest(n.left, p, ncp, !vsplit);
+            }
+        }
     }
 
 }