bundles/org.simantics.utils.datastructures/src/org/simantics/utils/datastructures/collections/QuadTree.java

   1 package org.simantics.utils.datastructures.collections;
   2
   3 import java.awt.geom.Point2D;
   4 import java.awt.geom.Rectangle2D;
   5 import java.util.HashSet;
   6 import java.util.Set;
   7
   8 public class QuadTree<T> {
   9
  10         Point2D center;
  11         Set<T> contains;
  12         double width;
  13         double height;
  14
  15         boolean leaf;
  16         QuadTree<T> pXpY;
  17         QuadTree<T> nXpY;
  18         QuadTree<T> pXnY;
  19         QuadTree<T> nXnY;
  20
  21         /**
  22          * Creates a quadtree
  23          * @param center center of the quadtree area
  24          * @param width width of the area
  25          * @param height height of the area
  26          * @param depth depth of the tree structure.
  27          */
  28         public QuadTree(Point2D center, double width, double height, int depth) {
  29                 this.center = center;
  30                 this.width = width;
  31                 this.height = height;
  32                 this.leaf = true;
  33                 split(depth);
  34         }
  35
  36         private QuadTree(Point2D center, double width, double height) {
  37                 this.center = center;
  38                 this.width = width;
  39                 this.height = height;
  40                 this.leaf = true;
  41         }
  42
  43         /**
  44          * Insets an object into the tree
  45          * @param object
  46          * @param bounds
  47          */
  48         public void insert(T object, Rectangle2D bounds) {
  49                 if (leaf) {
  50                         contains.add(object);
  51                 } else {
  52
  53                         boolean pX = bounds.getMinX() > center.getX();
  54                         boolean nX = bounds.getMaxX() < center.getX();
  55                         boolean pY = bounds.getMinY() > center.getY();
  56                         boolean nY = bounds.getMaxY() < center.getY();
  57
  58                         if (pX) {
  59                                 if (pY) {
  60                                         pXpY.insert(object, bounds);
  61                                 } else if (nY) {
  62                                         pXnY.insert(object, bounds);
  63                                 } else {
  64                                         pXpY.insert(object, bounds);
  65                                         pXnY.insert(object, bounds);
  66                                 }
  67                         } else if (nX) {
  68                                 if (pY) {
  69                                         nXpY.insert(object, bounds);
  70                                 } else if (nY) {
  71                                         nXnY.insert(object, bounds);
  72                                 } else {
  73                                         nXpY.insert(object, bounds);
  74                                         nXnY.insert(object, bounds);
  75                                 }
  76                         } else if (pY) {
  77                                 pXpY.insert(object, bounds);
  78                                 nXpY.insert(object, bounds);
  79                         } else if (nY) {
  80                                 pXnY.insert(object, bounds);
  81                                 nXnY.insert(object, bounds);
  82                         } else {
  83                                 pXpY.insert(object, bounds);
  84                                 pXnY.insert(object, bounds);
  85                                 nXpY.insert(object, bounds);
  86                                 nXnY.insert(object, bounds);
  87                         }
  88
  89                 }
  90         }
  91
  92         /**
  93          * Removes an object from the tree
  94          * @param object
  95          */
  96         public void remove(T object) {
  97                 if (leaf) {
  98                         contains.remove(object);
  99                 } else {
 100                         pXpY.remove(object);
 101                         pXnY.remove(object);
 102                         nXpY.remove(object);
 103                         nXnY.remove(object);
 104                 }
 105         }
 106
 107         /**
 108          * Returns objects within the given area.
 109          * @param bounds
 110          * @param set
 111          */
 112         public void get(Rectangle2D bounds, Set<T> set) {
 113                 if (leaf) {
 114                         set.addAll(contains);
 115                 } else {
 116
 117                         boolean pX = bounds.getMinX() > center.getX();
 118                         boolean nX = bounds.getMaxX() < center.getX();
 119                         boolean pY = bounds.getMinY() > center.getY();
 120                         boolean nY = bounds.getMaxY() < center.getY();
 121
 122                         if (pX) {
 123                                 if (pY) {
 124                                         pXpY.get(bounds, set);
 125                                 } else if (nY) {
 126                                         pXnY.get(bounds, set);
 127                                 } else {
 128                                         pXpY.get(bounds, set);
 129                                         pXnY.get(bounds, set);
 130                                 }
 131                         } else if (nX) {
 132                                 if (pY) {
 133                                         nXpY.get(bounds, set);
 134                                 } else if (nY) {
 135                                         nXnY.get(bounds, set);
 136                                 } else {
 137                                         nXpY.get(bounds, set);
 138                                         nXnY.get(bounds, set);
 139                                 }
 140                         } else if (pY) {
 141                                 pXpY.get(bounds, set);
 142                                 nXpY.get(bounds, set);
 143                         } else if (nY) {
 144                                 pXnY.get(bounds, set);
 145                                 nXnY.get(bounds, set);
 146                         } else {
 147                                 pXpY.get(bounds, set);
 148                                 pXnY.get(bounds, set);
 149                                 nXpY.get(bounds, set);
 150                                 nXnY.get(bounds, set);
 151                         }
 152
 153                 }
 154         }
 155
 156         public void clear() {
 157                 if (leaf) {
 158                         contains.clear();
 159                 } else {
 160                         pXpY.clear();
 161                         pXnY.clear();
 162                         nXpY.clear();
 163                         nXnY.clear();
 164                 }
 165         }
 166
 167         private void split(int depth) {
 168                 if (!leaf) {
 169                         throw new IllegalStateException("Node is already split");
 170                 }
 171                 if (depth <= 0) {
 172                         this.contains = new HashSet<>();
 173                         return;
 174                 }
 175                 split();
 176                 depth--;
 177                 pXpY.split(depth);
 178                 nXpY.split(depth);
 179                 pXnY.split(depth);
 180                 nXnY.split(depth);
 181         }
 182
 183         private void split() {
 184                 double w = width * 0.5;
 185                 double h = height * 0.5;
 186                 double wd2 = w * 0.5;
 187                 double hd2 = h * 0.5;
 188                 pXpY = new QuadTree<T>(new Point2D.Double(center.getX()+wd2, center.getY()+hd2), w, h);
 189                 nXpY = new QuadTree<T>(new Point2D.Double(center.getX()-wd2, center.getY()+hd2), w, h);
 190                 pXnY = new QuadTree<T>(new Point2D.Double(center.getX()+wd2, center.getY()-hd2), w, h);
 191                 nXnY = new QuadTree<T>(new Point2D.Double(center.getX()-wd2, center.getY()-hd2), w, h);
 192                 leaf = false;
 193         }
 194
 195
 196 }