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