1   package org.djutils.quadtree;
2   
3   import java.util.Collection;
4   import java.util.Iterator;
5   import java.util.LinkedHashSet;
6   import java.util.Set;
7   
8   import org.djutils.exceptions.Throw;
9   
10  
11  
12  
13  
14  
15  
16  
17  
18  
19  
20  public class QuadTree<T extends Envelope> implements Collection<T>
21  {
22      
23      private final int maximumLoad;
24  
25      
26      private final double minimumSize;
27  
28      
29      private final SubTree<T> tree;
30  
31      
32      private int totalSubTrees = 0;
33  
34      
35  
36  
37  
38  
39  
40  
41  
42  
43      public QuadTree(final int maximumLoad, final double minimumSize, final double left, final double bottom, final double right,
44              final double top)
45      {
46          Throw.when(left >= right, IllegalArgumentException.class, "left (%f) must be less than right (%f)", left, right);
47          Throw.when(bottom >= top, IllegalArgumentException.class, "bottom (%f) must be less than top (%f)", bottom, top);
48          this.maximumLoad = maximumLoad;
49          this.minimumSize = minimumSize;
50          this.tree = new SubTree<T>(this, new Rectangle(left, bottom, right, top));
51      }
52  
53      
54  
55  
56  
57      public int getMaxLoad()
58      {
59          return this.maximumLoad;
60      }
61  
62      
63  
64  
65  
66      public double getMinimumSize()
67      {
68          return this.minimumSize;
69      }
70  
71      @Override
72      public int size()
73      {
74          return this.tree.size();
75      }
76  
77      @Override
78      public boolean isEmpty()
79      {
80          return this.tree.size() == 0;
81      }
82  
83      @Override
84      public boolean contains(final Object o)
85      {
86          if (!(o instanceof Envelope))
87          {
88              return false;
89          }
90          @SuppressWarnings("unchecked")
91          T t = (T) o;
92          return this.tree.recursiveContains(new RectangleAndPayload<T>(t.getBoundingRectangle(), t));
93      }
94  
95      @Override
96      public Iterator<T> iterator()
97      {
98          return iterator(this.tree.getBoundingBox());
99      }
100 
101     
102 
103 
104 
105 
106 
107     public Iterator<T> iterator(final Rectangle searchArea)
108     {
109         return collect(searchArea).iterator();
110     }
111 
112     @Override
113     public Object[] toArray()
114     {
115         return collect(this.tree.getBoundingBox()).toArray();
116     }
117 
118     @SuppressWarnings("hiding")
119     @Override
120     public <T> T[] toArray(final T[] a)
121     {
122         return collect(this.tree.getBoundingBox()).toArray(a);
123     }
124 
125     
126 
127 
128 
129 
130     private Set<T> collect(final Rectangle searchArea)
131     {
132         Iterator<RectangleAndPayload<T>> iterator = this.tree.recursiveCollect(searchArea).iterator();
133         Set<T> result = new LinkedHashSet<>();
134         while (iterator.hasNext())
135         {
136             result.add(iterator.next().getPayload());
137         }
138         return result;
139     }
140 
141     @Override
142     public boolean add(final T e)
143     {
144         return this.tree.add(new RectangleAndPayload<T>(e.getBoundingRectangle(), e));
145     }
146 
147     @Override
148     public boolean remove(final Object o)
149     {
150         if (!(o instanceof Envelope))
151         {
152             return false;
153         }
154         @SuppressWarnings("unchecked")
155         T t = (T) o;
156         return this.tree.remove(new RectangleAndPayload<T>(t.getBoundingRectangle(), t));
157     }
158 
159     @Override
160     public boolean containsAll(final Collection<?> c)
161     {
162         return collect(this.tree.getBoundingBox()).containsAll(c);
163     }
164 
165     @Override
166     public boolean addAll(final Collection<? extends T> c)
167     {
168         boolean result = false;
169         for (T t : c)
170         {
171             result |= add(t);
172         }
173         return result;
174     }
175 
176     @Override
177     public boolean removeAll(final Collection<?> c)
178     {
179         boolean result = false;
180         for (Object o : c)
181         {
182             result |= remove(o);
183         }
184         return result;
185     }
186 
187     @Override
188     public boolean retainAll(final Collection<?> c)
189     {
190         throw new RuntimeException("Not (yet) implemented");
191     }
192 
193     @Override
194     public void clear()
195     {
196         this.tree.clear();
197     }
198 
199     
200 
201 
202     void incrementSubTreeCount()
203     {
204         this.totalSubTrees++;
205     }
206 
207     
208 
209 
210 
211     public int getSubTreeCount()
212     {
213         return this.totalSubTrees;
214     }
215 
216     @Override
217     public String toString()
218     {
219         return "QuadTree [maximumLoad=" + this.maximumLoad + ", minimumSize=" + this.minimumSize + ", tree=" + this.tree + "]";
220     }
221 
222     
223 
224 
225 
226 
227     public String toString(final int expandDepth)
228     {
229         return "QuadTree [maximumLoad=" + this.maximumLoad + ", minimumSize=" + this.minimumSize + ", tree="
230                 + this.tree.toString(expandDepth) + "]";
231     }
232 
233     
234 
235 
236 
237 
238     public String dump(final String indent)
239     {
240         return this.tree.dump(indent);
241     }
242 
243     
244 
245 
246 
247     @SuppressWarnings("hiding")
248     class SubTree<T extends Envelope>
249     {
250         
251         private final QuadTree<T> root;
252 
253         
254         private final Rectangle boundingBox;
255 
256         
257         private int size = 0;
258 
259         
260         private SubTree<T>[] children = null;
261 
262         
263         private Set<RectangleAndPayload<T>> elements = null;
264 
265         
266 
267 
268 
269 
270         SubTree(final QuadTree<T> root, final Rectangle boundingBox)
271         {
272             this.root = root;
273             this.boundingBox = boundingBox;
274             root.incrementSubTreeCount();
275         }
276 
277         
278 
279 
280 
281         public final Rectangle getBoundingBox()
282         {
283             return this.boundingBox;
284         }
285 
286         
287 
288 
289 
290         public int size()
291         {
292             return this.size;
293         }
294 
295         
296 
297 
298 
299 
300         public boolean add(final RectangleAndPayload<T> e)
301         {
302             if (contains(e))
303             {
304                 return false;
305             }
306             if (this.elements == null)
307             {
308                 this.elements = new LinkedHashSet<>();
309             }
310             this.elements.add(e);
311             this.size++;
312             reBalance();
313             return true;
314         }
315 
316         
317 
318 
319 
320 
321         public boolean remove(final RectangleAndPayload<T> o)
322         {
323             if (this.elements.remove(o))
324             {
325                 this.size--;
326                 return true; 
327             }
328             
329             boolean result = false;
330             if (this.children != null)
331             {
332                 Rectangle rectangle = o.getRectangle();
333                 for (SubTree<T> child : this.children)
334                 {
335                     if (!child.boundingBox.intersects(rectangle))
336                     {
337                         continue; 
338                     }
339                     if (child.remove(o))
340                     {
341                         result = true;
342                     }
343                 }
344             }
345             if (result)
346             {
347                 this.size--;
348             }
349             return result;
350         }
351 
352         
353 
354 
355         public void clear()
356         {
357             this.elements.clear();
358             this.children = null;
359             this.size = 0;
360         }
361 
362         
363 
364 
365 
366 
367         public boolean contains(final RectangleAndPayload<T> o)
368         {
369             if (this.elements == null)
370             {
371                 return false;
372             }
373             return recursiveContains(o);
374         }
375 
376         
377 
378 
379 
380 
381         boolean recursiveContains(final RectangleAndPayload<T> o)
382         {
383             if ((!this.boundingBox.intersects(o.getRectangle())) || (this.elements == null))
384             {
385                 return false; 
386             }
387             for (RectangleAndPayload<T> element : this.elements)
388             {
389                 if (element.equals(o))
390                 {
391                     return true;
392                 }
393             }
394             if (this.children == null)
395             {
396                 return false;
397             }
398             for (SubTree<T> child : this.children)
399             {
400                 if (child.recursiveContains(o))
401                 {
402                     return true;
403                 }
404             }
405             return false;
406         }
407 
408         
409 
410 
411 
412 
413         public Set<RectangleAndPayload<T>> recursiveCollect(final Rectangle rectangle)
414         {
415             Set<RectangleAndPayload<T>> result = new LinkedHashSet<>();
416             if (!this.boundingBox.intersects(rectangle))
417             {
418                 return result; 
419             }
420             if (this.elements != null)
421             {
422                 for (RectangleAndPayload<T> element : this.elements)
423                 {
424                     if (element.getRectangle().intersects(rectangle))
425                     {
426                         result.add(element);
427                     }
428                 }
429             }
430             if (this.children != null)
431             {
432                 for (SubTree<T> child : this.children)
433                 {
434                     result.addAll(child.recursiveCollect(rectangle));
435                 }
436             }
437             return result;
438         }
439 
440         
441 
442 
443         @SuppressWarnings("unchecked")
444         private void reBalance()
445         {
446             if (this.elements.size() < this.root.getMaxLoad() || this.boundingBox.getWidth() < this.root.getMinimumSize()
447                     || this.boundingBox.getHeight() < this.root.getMinimumSize())
448             {
449                 return;
450             }
451             
452             double cX = (this.boundingBox.getLeft() + this.boundingBox.getRight()) / 2;
453             double cY = (this.boundingBox.getBottom() + this.boundingBox.getTop()) / 2;
454             int canMove = 0;
455             
456 
457 
458 
459 
460 
461 
462 
463 
464 
465             canMove = this.elements.size();
466             if (canMove == 0 || canMove < this.root.getMaxLoad()  && this.children == null)
467             {
468                 
469                 return;
470             }
471             
472             if (this.children == null)
473             {
474                 this.children = new SubTree[] {
475                         new SubTree<T>(this.root,
476                                 new Rectangle(this.boundingBox.getLeft(), this.boundingBox.getBottom(), cX, cY)),
477                         new SubTree<T>(this.root,
478                                 new Rectangle(cX, this.boundingBox.getBottom(), this.boundingBox.getRight(), cY)),
479                         new SubTree<T>(this.root, new Rectangle(this.boundingBox.getLeft(), cY, cX, this.boundingBox.getTop())),
480                         new SubTree<T>(this.root,
481                                 new Rectangle(cX, cY, this.boundingBox.getRight(), this.boundingBox.getTop()))};
482             }
483             Iterator<RectangleAndPayload<T>> iterator = this.elements.iterator();
484             while (iterator.hasNext())
485             {
486                 RectangleAndPayload<T> e = iterator.next();
487                 if (e.getRectangle().contains(this.boundingBox))
488                 {
489                     continue;
490                 }
491                 boolean added = false;
492                 for (SubTree<T> child : this.children)
493                 {
494                     if (e.getRectangle().intersects(child.boundingBox))
495                     {
496                         added |= child.add(e);
497                     }
498                 }
499                 if (added)
500                 {
501                     iterator.remove();
502                 }
503                 else
504                 {
505                     System.out.println("ERROR: Could not add " + e + " to any of the children");
506                 }
507             }
508             
509         }
510 
511         @Override
512         public String toString()
513         {
514             return "SubTree [boundingBox=" + this.boundingBox + ", size=" + this.size + ", children=" + this.children
515                     + (this.elements == null ? "elements=null" : ", elements.size=" + this.elements.size()) + "]";
516         }
517 
518         
519 
520 
521 
522 
523         public String toString(final int expandDepth)
524         {
525             if (expandDepth > 0)
526             {
527                 return "SubTree [boundingBox=" + this.boundingBox + ", size=" + this.size + ", children="
528                         + (this.children != null ? "[SW:" + this.children[0].toString(expandDepth - 1) + ", SE:"
529                                 + this.children[1].toString(expandDepth - 1) + ", NW:"
530                                 + this.children[2].toString(expandDepth - 1) + ", NE:"
531                                 + this.children[3].toString(expandDepth - 1) + "]" : "null")
532                         + ", elements.size=" + this.elements.size() + "]";
533             }
534             else
535             {
536                 return toString();
537             }
538         }
539 
540         
541 
542 
543 
544 
545         public String dump(final String indent)
546         {
547             StringBuilder result = new StringBuilder();
548             result.append(indent);
549             result.append("SubTree [size=");
550             result.append(this.size);
551             result.append("] ");
552             result.append(this.boundingBox);
553             result.append("\n");
554             String subIndent = indent + "    ";
555             Iterator<RectangleAndPayload<T>> iterator = this.elements.iterator();
556             for (int i = 0; i < this.elements.size(); i++)
557             {
558                 result.append(subIndent);
559                 result.append(i);
560                 result.append(" ");
561                 result.append(iterator.next());
562                 result.append("\n");
563             }
564             if (this.children != null)
565             {
566                 result.append(subIndent);
567                 result.append("SW");
568                 result.append("\n");
569                 result.append(this.children[0].dump(subIndent));
570                 result.append(subIndent);
571                 result.append("SE");
572                 result.append("\n");
573                 result.append(this.children[1].dump(subIndent));
574                 result.append(subIndent);
575                 result.append("NW");
576                 result.append("\n");
577                 result.append(this.children[2].dump(subIndent));
578                 result.append(subIndent);
579                 result.append("NE");
580                 result.append("\n");
581                 result.append(this.children[3].dump(subIndent));
582             }
583             return result.toString();
584         }
585 
586     }
587 
588 }
589 
590 
591 
592 
593 
594 class RectangleAndPayload<T extends Object>
595 {
596     
597     private final Rectangle rectangle;
598 
599     
600     private final T payload;
601 
602     
603 
604 
605 
606 
607     RectangleAndPayload(final Rectangle rectangle, final T payload)
608     {
609         this.rectangle = rectangle;
610         this.payload = payload;
611     }
612 
613     
614 
615 
616 
617     public Rectangle getRectangle()
618     {
619         return this.rectangle;
620     }
621 
622     
623 
624 
625 
626     public T getPayload()
627     {
628         return this.payload;
629     }
630 
631     @Override
632     public String toString()
633     {
634         return "RectangleAndPayload [rectangle=" + this.rectangle + ", payload=" + this.payload + "]";
635     }
636 
637     @Override
638     public int hashCode()
639     {
640         final int prime = 31;
641         int result = 1;
642         result = prime * result + ((this.payload == null) ? 0 : this.payload.hashCode());
643         result = prime * result + ((this.rectangle == null) ? 0 : this.rectangle.hashCode());
644         return result;
645     }
646 
647     @Override
648     public boolean equals(final Object obj)
649     {
650         if (this == obj)
651         {
652             return true;
653         }
654         if (obj == null)
655         {
656             return false;
657         }
658         if (getClass() != obj.getClass())
659         {
660             return false;
661         }
662         @SuppressWarnings("rawtypes")
663         RectangleAndPayload other = (RectangleAndPayload) obj;
664         if (this.payload == null)
665         {
666             if (other.payload != null)
667             {
668                 return false;
669             }
670         }
671         else if (!this.payload.equals(other.payload))
672         {
673             return false;
674         }
675         if (this.rectangle == null)
676         {
677             if (other.rectangle != null)
678             {
679                 return false;
680             }
681         }
682         else if (!this.rectangle.equals(other.rectangle))
683         {
684             return false;
685         }
686         return true;
687     }
688 
689 }