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

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