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 }