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         public void remove(T object, Rectangle2D bounds) {
 108                 if (leaf) {
 109                         contains.remove(object);
 110                 } else {
 111
 112                         boolean pX = bounds.getMinX() > center.getX();
 113                         boolean nX = bounds.getMaxX() < center.getX();
 114                         boolean pY = bounds.getMinY() > center.getY();
 115                         boolean nY = bounds.getMaxY() < center.getY();
 116
 117                         if (pX) {
 118                                 if (pY) {
 119                                         pXpY.remove(object, bounds);
 120                                 } else if (nY) {
 121                                         pXnY.remove(object, bounds);
 122                                 } else {
 123                                         pXpY.remove(object, bounds);
 124                                         pXnY.remove(object, bounds);
 125                                 }
 126                         } else if (nX) {
 127                                 if (pY) {
 128                                         nXpY.remove(object, bounds);
 129                                 } else if (nY) {
 130                                         nXnY.remove(object, bounds);
 131                                 } else {
 132                                         nXpY.remove(object, bounds);
 133                                         nXnY.remove(object, bounds);
 134                                 }
 135                         } else if (pY) {
 136                                 pXpY.remove(object, bounds);
 137                                 nXpY.remove(object, bounds);
 138                         } else if (nY) {
 139                                 pXnY.remove(object, bounds);
 140                                 nXnY.remove(object, bounds);
 141                         } else {
 142                                 pXpY.remove(object, bounds);
 143                                 pXnY.remove(object, bounds);
 144                                 nXpY.remove(object, bounds);
 145                                 nXnY.remove(object, bounds);
 146                         }
 147
 148                 }
 149         }
 150
 151         public boolean contains(T object) {
 152                 if (leaf) {
 153                         return contains.contains(object);
 154                 } else {
 155                         return pXnY.contains(object) || pXpY.contains(object) || nXnY.contains(object) || nXpY.contains(object);
 156                 }
 157
 158         }
 159
 160         /**
 161          * Returns objects within the given area.
 162          * @param bounds
 163          * @param set
 164          */
 165         public void get(Rectangle2D bounds, Set<T> set) {
 166                 if (leaf) {
 167                         set.addAll(contains);
 168                 } else {
 169
 170                         boolean pX = bounds.getMinX() > center.getX();
 171                         boolean nX = bounds.getMaxX() < center.getX();
 172                         boolean pY = bounds.getMinY() > center.getY();
 173                         boolean nY = bounds.getMaxY() < center.getY();
 174
 175                         if (pX) {
 176                                 if (pY) {
 177                                         pXpY.get(bounds, set);
 178                                 } else if (nY) {
 179                                         pXnY.get(bounds, set);
 180                                 } else {
 181                                         pXpY.get(bounds, set);
 182                                         pXnY.get(bounds, set);
 183                                 }
 184                         } else if (nX) {
 185                                 if (pY) {
 186                                         nXpY.get(bounds, set);
 187                                 } else if (nY) {
 188                                         nXnY.get(bounds, set);
 189                                 } else {
 190                                         nXpY.get(bounds, set);
 191                                         nXnY.get(bounds, set);
 192                                 }
 193                         } else if (pY) {
 194                                 pXpY.get(bounds, set);
 195                                 nXpY.get(bounds, set);
 196                         } else if (nY) {
 197                                 pXnY.get(bounds, set);
 198                                 nXnY.get(bounds, set);
 199                         } else {
 200                                 pXpY.get(bounds, set);
 201                                 pXnY.get(bounds, set);
 202                                 nXpY.get(bounds, set);
 203                                 nXnY.get(bounds, set);
 204                         }
 205
 206                 }
 207         }
 208
 209         public void clear() {
 210                 if (leaf) {
 211                         contains.clear();
 212                 } else {
 213                         pXpY.clear();
 214                         pXnY.clear();
 215                         nXpY.clear();
 216                         nXnY.clear();
 217                 }
 218         }
 219
 220         private void split(int depth) {
 221                 if (!leaf) {
 222                         throw new IllegalStateException("Node is already split");
 223                 }
 224                 if (depth <= 0) {
 225                         this.contains = new HashSet<>();
 226                         return;
 227                 }
 228                 split();
 229                 depth--;
 230                 pXpY.split(depth);
 231                 nXpY.split(depth);
 232                 pXnY.split(depth);
 233                 nXnY.split(depth);
 234         }
 235
 236         private void split() {
 237                 double w = width * 0.5;
 238                 double h = height * 0.5;
 239                 double wd2 = w * 0.5;
 240                 double hd2 = h * 0.5;
 241                 pXpY = new QuadTree<T>(new Point2D.Double(center.getX()+wd2, center.getY()+hd2), w, h);
 242                 nXpY = new QuadTree<T>(new Point2D.Double(center.getX()-wd2, center.getY()+hd2), w, h);
 243                 pXnY = new QuadTree<T>(new Point2D.Double(center.getX()+wd2, center.getY()-hd2), w, h);
 244                 nXnY = new QuadTree<T>(new Point2D.Double(center.getX()-wd2, center.getY()-hd2), w, h);
 245                 leaf = false;
 246         }
 247
 248
 249 }