quadtree.js

// Daniel Shiffman
// http://codingtra.in
// http://patreon.com/codingtrain

// QuadTree

class Point {
  constructor(x, y, data) {
    this.x = x;
    this.y = y;
    this.userData = data;
  }
}

class Rectangle {
  constructor(x, y, w, h) {
    this.x = x;
    this.y = y;
    this.w = w;
    this.h = h;
  }

  contains(point) {
    return (point.x >= this.x - this.w &&
      point.x <= this.x + this.w &&
      point.y >= this.y - this.h &&
      point.y <= this.y + this.h);
  }


  intersects(range) {
    return !(range.x - range.w > this.x + this.w ||
      range.x + range.w < this.x - this.w ||
      range.y - range.h > this.y + this.h ||
      range.y + range.h < this.y - this.h);
  }


}

// circle class for a circle shaped query
class Circle {
  constructor(x, y, r) {
    this.x = x;
    this.y = y;
    this.r = r;
    this.rSquared = this.r * this.r;
  }

  contains(point) {
    // check if the point is in the circle by checking if the euclidean distance of
    // the point and the center of the circle if smaller or equal to the radius of
    // the circle
    let d = Math.pow((point.x - this.x), 2) + Math.pow((point.y - this.y), 2);
    return d <= this.rSquared;
  }

  intersects(range) {

    let xDist = Math.abs(range.x - this.x);
    let yDist = Math.abs(range.y - this.y);

    // radius of the circle
    let r = this.r;

    let w = range.w;
    let h = range.h;

    let edges = Math.pow((xDist - w), 2) + Math.pow((yDist - h), 2);

    // no intersection
    if (xDist > (r + w) || yDist > (r + h))
      return false;

    // intersection within the circle
    if (xDist <= w || yDist <= h)
      return true;

    // intersection on the edge of the circle
    return edges <= this.rSquared;
  }
}

class QuadTree {
  constructor(boundary, capacity) {
    if (!boundary) {
      throw TypeError('boundary is null or undefined');
    }
    if (!(boundary instanceof Rectangle)) {
      throw TypeError('boundary should be a Rectangle');
    }
    if (typeof capacity !== 'number') {
      throw TypeError(`capacity should be a number but is a ${typeof capacity}`);
    }
    if (capacity < 1) {
      throw RangeError('capacity must be greater than 0');
    }
    this.boundary = boundary;
    this.capacity = capacity;
    this.points = [];
    this.divided = false;
  }

  toJSON(isChild) {
    let obj = { points: this.points };
    if (this.divided) {
      if (this.northeast.points.length > 0) {
        obj.ne = this.northeast.toJSON(true);
      }
      if (this.northwest.points.length > 0) {
        obj.nw = this.northwest.toJSON(true);
      }
      if (this.southeast.points.length > 0) {
        obj.se = this.southeast.toJSON(true);
      }
      if (this.southwest.points.length > 0) {
        obj.sw = this.southwest.toJSON(true);
      }
    }
    if (!isChild) {
      obj.capacity = this.capacity;
      obj.x = this.boundary.x;
      obj.y = this.boundary.y;
      obj.w = this.boundary.w;
      obj.h = this.boundary.h;
    }
    return obj;
  }

  static fromJSON(obj, x, y, w, h, capacity) {
    if (typeof x === "undefined") {
      if ("x" in obj) {
        x = obj.x;
        y = obj.y;
        w = obj.w;
        h = obj.h;
        capacity = obj.capacity;
      } else {
        throw TypeError("JSON missing boundary information");
      }
    }
    let qt = new QuadTree(new Rectangle(x, y, w, h), capacity);
    qt.points = obj.points;
    if (
      "ne" in obj ||
      "nw" in obj ||
      "se" in obj ||
      "sw" in obj
    ) {
      let x = qt.boundary.x;
      let y = qt.boundary.y;
      let w = qt.boundary.w / 2;
      let h = qt.boundary.h / 2;

      if ("ne" in obj) {
        qt.northeast = QuadTree.fromJSON(obj.ne, x + w, y - h, w, h, capacity);
      } else {
        qt.northeast = new QuadTree(new Rectangle(x + w, y - h, w, h), capacity);
      }
      if ("nw" in obj) {
        qt.northwest = QuadTree.fromJSON(obj.nw, x - w, y - h, w, h, capacity);
      } else {
        qt.northwest = new QuadTree(new Rectangle(x - w, y - h, w, h), capacity);
      }
      if ("se" in obj) {
        qt.southeast = QuadTree.fromJSON(obj.se, x + w, y + h, w, h, capacity);
      } else {
        qt.southeast = new QuadTree(new Rectangle(x + w, y + h, w, h), capacity);
      }
      if ("sw" in obj) {
        qt.southwest = QuadTree.fromJSON(obj.sw, x - w, y + h, w, h, capacity);
      } else {
        qt.southwest = new QuadTree(new Rectangle(x - w, y + h, w, h), capacity);
      }

      qt.divided = true;
    }
    return qt;
  }

  subdivide() {
    let x = this.boundary.x;
    let y = this.boundary.y;
    let w = this.boundary.w / 2;
    let h = this.boundary.h / 2;

    let ne = new Rectangle(x + w, y - h, w, h);
    this.northeast = new QuadTree(ne, this.capacity);
    let nw = new Rectangle(x - w, y - h, w, h);
    this.northwest = new QuadTree(nw, this.capacity);
    let se = new Rectangle(x + w, y + h, w, h);
    this.southeast = new QuadTree(se, this.capacity);
    let sw = new Rectangle(x - w, y + h, w, h);
    this.southwest = new QuadTree(sw, this.capacity);

    this.divided = true;
  }

  insert(point) {
    if (!this.boundary.contains(point)) {
      return false;
    }

    if (this.points.length < this.capacity) {
      this.points.push(point);
      return true;
    }

    if (!this.divided) {
      this.subdivide();
    }

    return (this.northeast.insert(point) || this.northwest.insert(point) ||
      this.southeast.insert(point) || this.southwest.insert(point));
  }

  query(range, found) {
    if (!found) {
      found = [];
    }

    if (!range.intersects(this.boundary)) {
      return found;
    }

    for (let p of this.points) {
      if (range.contains(p)) {
        found.push(p);
      }
    }
    if (this.divided) {
      this.northwest.query(range, found);
      this.northeast.query(range, found);
      this.southwest.query(range, found);
      this.southeast.query(range, found);
    }

    return found;
  }

  closest(point, count, startingSize) {
    if (typeof count === "undefined") {
      count = 1;
    }
    if (typeof startingSize === "undefined") {
      startingSize = 1;
    }

    if (!this.divided) {
      // Empty
      if (this.points.length == 0) {
        return [];
      }
      // Limit to number of points in this QuadTree
      if (this.points.length < count) {
        count = this.points.length;
      }
    }

    // optimized, expanding binary search
    // start with a small circle, rapidly expand, slowly shrink
    let radius = startingSize;
    let limit = 16;
    while (true) {
      const range = new Circle(point.x, point.y, radius);
      const points = this.query(range);
      if (points.length === count) {
        return points; // Return when we hit the right size
      } else if (points.length < count) {
        radius *= 2;
      } else if (limit-- <= 0) {
        return points.slice(0, count); // Slice to return correct count (breaks ties)
      } else {
        radius /= 3;
      }
    }
  }
}

if (typeof module !== "undefined") {
  module.exports = { Point, Rectangle, QuadTree, Circle };
}