算法导论第四版学习——习题五Kd-Tree

 1 import edu.princeton.cs.algs4.Point2D;
 2 import edu.princeton.cs.algs4.RectHV;
 3 import edu.princeton.cs.algs4.SET;
 4 
 5 import java.util.Stack;
 6 
 7 
 8 public class PointSET {
 9     private SET<Point2D> pointSet;
10 
11     public PointSET() // construct an empty set of points 
12      {
13         pointSet = new SET<Point2D>();
14     }
15 
16     private void checkNull(Object obj) {
17         if (obj == null) {
18             throw new NullPointerException();
19         }
20     }
21 
22     public boolean isEmpty() // is the set empty? 
23      {
24         return pointSet.isEmpty();
25     }
26 
27     public int size() // number of points in the set 
28      {
29         return pointSet.size();
30     }
31 
32     public void insert(Point2D p) // add the point to the set (if it is not already in the set)
33      {
34         checkNull(p);
35         pointSet.add(p);
36     }
37 
38     public boolean contains(Point2D p) // does the set contain point p? 
39      {
40         checkNull(p);
41 
42         return pointSet.contains(p);
43     }
44 
45     public void draw() // draw all points to standard draw 
46      {
47         for (Point2D p : pointSet) {
48             p.draw();
49         }
50     }
51 
52     public Iterable<Point2D> range(RectHV rect) // all points that are inside the rectangle 
53      {
54         checkNull(rect);
55 
56         Stack<Point2D> stack = new Stack<Point2D>();
57 
58         for (Point2D p : pointSet) {
59             if ((p.x() >= rect.xmin()) && (p.x() <= rect.xmax()) &&
60                     (p.y() >= rect.ymin()) && (p.y() <= rect.ymax())) {
61                 stack.push(p);
62             }
63         }
64 
65         return stack;
66     }
67 
68     public Point2D nearest(Point2D p) // a nearest neighbor in the set to point p; null if the set is empty 
69      {
70         checkNull(p);
71 
72         Point2D point = null;
73 
74         for (Point2D setp : pointSet) {
75             if ((point == null) ||
76                     (p.distanceSquaredTo(point) > p.distanceSquaredTo(setp))) {
77                 point = setp;
78             }
79         }
80 
81         return point;
82     }
83 
84     public static void main(String[] args) // unit testing of the methods (optional) 
85      {
86     }
87 }

  1 import edu.princeton.cs.algs4.In;
  2 import edu.princeton.cs.algs4.Point2D;
  3 import edu.princeton.cs.algs4.RectHV;
  4 import edu.princeton.cs.algs4.StdDraw;
  5 import edu.princeton.cs.algs4.StdOut;
  6 
  7 import java.util.Stack;
  8 
  9 
 10 public class KdTree {
 11     private PointNode root;
 12     private int pointCount;
 13 
 14     public KdTree() // construct an empty set of points 
 15      {
 16         root = null;
 17         pointCount = 0;
 18     }
 19 
 20     private void checkNull(Object obj) {
 21         if (obj == null) {
 22             throw new NullPointerException();
 23         }
 24     }
 25 
 26     public boolean isEmpty() // is the set empty? 
 27      {
 28         return root == null;
 29     }
 30 
 31     public int size() // number of points in the set 
 32      {
 33         return pointCount;
 34     }
 35 
 36     public void insert(Point2D p) // add the point to the set (if it is not already in the set)
 37      {
 38         checkNull(p);
 39 
 40         if (root == null) {
 41             root = new PointNode(p, null, false);
 42             pointCount++;
 43         } else {
 44             PointNode current = root;
 45 
 46             while (true) {
 47               if (current.value.compareTo(p) == 0) {
 48                 return;
 49               }
 50                 if (current.isAtLeftSideOfNode(p)) {
 51                     if (current.left == null) {
 52                         current.left = new PointNode(p, current, true);
 53                         pointCount++;
 54                         break;
 55                     } else {
 56                         current = current.left;
 57                     }
 58                 } else {
 59                     if (current.right == null) {
 60                         current.right = new PointNode(p, current, false);
 61                         pointCount++;
 62                         break;
 63                     } else {
 64                         current = current.right;
 65                     }
 66                 }
 67             }
 68         }
 69     }
 70 
 71     public boolean contains(Point2D p) // does the set contain point p? 
 72      {
 73         checkNull(p);
 74 
 75         PointNode current = root;
 76 
 77         while (current != null) {
 78             if (current.value.compareTo(p) == 0) {
 79                 return true;
 80             } else if (current.isAtLeftSideOfNode(p)) {
 81                 current = current.left;
 82             } else {
 83                 current = current.right;
 84             }
 85         }
 86 
 87         return false;
 88     }
 89 
 90     public void draw() // draw all points to standard draw 
 91      {
 92         if (root != null) {
 93             root.draw(null, false);
 94         }
 95     }
 96 
 97     private void addToStack(Stack<Point2D> stack, PointNode current, RectHV rect) {
 98         
 99         if (!rect.intersects(current.rect)) {
100             return;
101         }
102 
103         if ((current.value.x() >= rect.xmin()) &&
104                 (current.value.x() <= rect.xmax()) &&
105                 (current.value.y() >= rect.ymin()) &&
106                 (current.value.y() <= rect.ymax())) {
107             stack.push(current.value);
108         }
109 
110         if (current.left != null) {
111             addToStack(stack, current.left, rect);
112         }
113 
114         if (current.right != null) {
115             addToStack(stack, current.right, rect);
116         }
117     }
118 
119     private Point2D searchNode(Point2D toSearch, PointNode current,
120         Point2D nearestPoint) {
121         if (nearestPoint == null) {
122             nearestPoint = current.value;
123         }
124 
125         double distance = nearestPoint.distanceSquaredTo(toSearch);
126         double newdistance = current.value.distanceSquaredTo(toSearch);
127 
128         if ((distance >= current.rect.distanceSquaredTo(toSearch)) ||
129                 (distance >= newdistance)) {
130             if (distance > newdistance) {
131                 nearestPoint = current.value;
132             }
133 
134             if ((current.isH && (toSearch.y() > current.value.y())) ||
135                     (!current.isH && (toSearch.x() < current.value.x()))) {
136                 if (current.left != null) {
137                     nearestPoint = searchNode(toSearch, current.left,
138                             nearestPoint);
139                 }
140 
141                 if (current.right != null) {
142                     nearestPoint = searchNode(toSearch, current.right,
143                             nearestPoint);
144                 }
145             } else {
146                 if (current.right != null) {
147                     nearestPoint = searchNode(toSearch, current.right,
148                             nearestPoint);
149                 }
150 
151                 if (current.left != null) {
152                     nearestPoint = searchNode(toSearch, current.left,
153                             nearestPoint);
154                 }
155             }
156         }
157 
158         return nearestPoint;
159     }
160 
161     public Iterable<Point2D> range(RectHV rect) // all points that are inside the rectangle 
162      {
163         checkNull(rect);
164 
165         Stack<Point2D> stack = new Stack<Point2D>();
166         if (root != null)
167           addToStack(stack, root, rect);
168 
169         return stack;
170     }
171 
172     public Point2D nearest(Point2D p) // a nearest neighbor in the set to point p; null if the set is empty 
173      {
174         checkNull(p);
175 
176         Point2D point = null;
177         if (root != null)
178           point = searchNode(p, root, point);
179 
180         return point;
181     }
182 
183     public static void main(String[] args) // unit testing of the methods (optional) 
184      {
185         String filename = args[0];
186         In in = new In(filename);
187 
188         KdTree kdtree = new KdTree();
189 
190         while (!in.isEmpty()) {
191             double x = in.readDouble();
192             double y = in.readDouble();
193             Point2D p = new Point2D(x, y);
194             kdtree.insert(p);
195         }
196 
197         kdtree.draw();
198     }
199 
200     private static class PointNode {
201         private PointNode left;
202         private PointNode right;
203         private RectHV rect;
204         private Point2D value;
205         private boolean isH;
206 
207         public PointNode(Point2D p, PointNode parent, boolean isLeftUp) {
208             value = p;
209             left = null;
210             right = null;
211             if (parent == null) {
212                 isH = false;
213                 rect = new RectHV(0, 0, 1, 1);
214             } else {
215                 isH = !parent.isH;
216 
217                 if (isH) {
218                     if (isLeftUp) {
219                         rect = new RectHV(parent.rect.xmin(),
220                                 parent.rect.ymin(), parent.value.x(),
221                                 parent.rect.ymax());
222                     } else {
223                         rect = new RectHV(parent.value.x(), parent.rect.ymin(),
224                                 parent.rect.xmax(), parent.rect.ymax());
225                     }
226                 } else {
227                     if (isLeftUp) {
228                         rect = new RectHV(parent.rect.xmin(), parent.value.y(),
229                                 parent.rect.xmax(), parent.rect.ymax());
230                     } else {
231                         rect = new RectHV(parent.rect.xmin(),
232                                 parent.rect.ymin(), parent.rect.xmax(),
233                                 parent.value.y());
234                     }
235                 }
236             }
237         }
238 
239         public void draw(PointNode parent, boolean isLeftUp) {
240             StdOut.println(isLeftUp ? "leftup" : "rightdown");
241             StdDraw.setPenColor(StdDraw.BLACK);
242             StdDraw.setPenRadius(0.01);
243             value.draw();
244             StdDraw.setPenRadius();
245 
246             if (parent == null) {
247                 StdDraw.setPenColor(StdDraw.RED);
248                 new Point2D(value.x(), 1).drawTo(new Point2D(value.x(), 0));
249             } else {
250                 StdOut.println(value);
251                 StdOut.println(parent.rect);
252 
253                 if (parent.isH) {
254                     StdDraw.setPenColor(StdDraw.RED);
255                     new Point2D(value.x(), rect.ymin()).drawTo(new Point2D(
256                             value.x(), rect.ymax()));
257                 } else {
258                     StdDraw.setPenColor(StdDraw.BLUE);
259                     new Point2D(rect.xmax(), value.y()).drawTo(new Point2D(
260                             rect.xmin(), value.y()));
261                 }
262             }
263 
264             StdDraw.pause(100);
265 
266             if (left != null) {
267                 left.draw(this, true);
268             }
269 
270             if (right != null) {
271                 right.draw(this, false);
272             }
273         }
274 
275         public boolean isAtLeftSideOfNode(Point2D p) {
276             if (isH) {
277                 return p.y() > value.y();
278             } else {
279                 return p.x() < value.x();
280             }
281         }
282     }
283 }