Quad Tree in ml

quadtree.ml

module Quadtree : QUADTREE = struct

  type 'a t = Rect.t * 'a qnode
  and 'a qnode = QLeaf of Pt.t * 'a list | QNode of 'a quadrants
  and 'a quadrants = 'a t * 'a t * 'a t * 'a t (* nw, ne, sw, se *)

  exception Out_of_bounds

  let make ((p, _) as r) = (r, QLeaf (p, []))

  let contains ((r, _) : 'a t) (p : Pt.t) : bool = Rect.contains r p

  let rec insert' ((r, qn) as qt : 'a t) (v : 'a) (p : Pt.t) : 'a t =
    if not (contains qt p) then raise Out_of_bounds else
    match qn with
    | QLeaf (p1, data) ->
      if data = [] || p = p1 then (r, QLeaf (p, v :: data))
      else
        let qt = (r, QNode (make (Rect.nw_quad r), make (Rect.ne_quad r),
                            make (Rect.sw_quad r), make (Rect.se_quad r))) in
        let qt = List.fold_left (fun qt d -> insert' qt d p1) qt data in
        insert' qt v p
    | QNode (nw, ne, sw, se) ->
      let quadrants =
        if contains nw p then (insert' nw v p, ne, sw, se)
        else if contains ne p then (nw, insert' ne v p, sw, se)
        else if contains sw p then (nw, ne, insert' sw v p, se)
        else if contains se p then (nw, ne, sw, insert' se v p)
        else raise Out_of_bounds in
      (r, QNode quadrants)

  (* type of insert should have been 'a t -> 'a -> Pt.t -> 'a t *)
  let insert qt v x y = insert' qt v (x, y)

  let rec fold_rect (f : Pt.t -> 'a -> 'b -> 'b) (accum : 'b) ((r, qn): 'a t) (rect : Rect.t) : 'b =
    if not (Rect.intersects r rect) then accum else
    match qn with
    | QLeaf (p, data) ->
      if data <> [] && Rect.contains rect p
      then List.fold_left (fun acc d -> f p d acc) accum data
      else accum
    | QNode (nw, ne, sw, se) ->
      let accum = fold_rect f accum nw rect in
      let accum = fold_rect f accum ne rect in
      let accum = fold_rect f accum sw rect in
      fold_rect f accum se rect

  type circle = Pt.t * float (* center, radius *)

  let circle_contains ((center, radius) : circle) (p : Pt.t) : bool =
    Pt.dist center p <= radius

  let circle_intersects ((center, radius) as c : circle) (((x1, y1), (x2, y2)) : Rect.t) : bool =
    circle_contains c (x1, y2) || circle_contains c (x2, y2) ||
    circle_contains c (x1, y1) || circle_contains c (x2, y1) ||
    Rect.contains ((x1 -. radius, y1), (x2 +. radius, y2)) center ||
    Rect.contains ((x1, y1 -. radius), (x2, y2 +. radius)) center

  (* very similar to fold_rect *)
  let rec fold_circ (f : Pt.t -> 'a -> 'b -> 'b) (accum : 'b) ((r, qn): 'a t) (c : circle) : 'b =
    if not (circle_intersects c r) then accum else
    match qn with
    | QLeaf (p, data) ->
      if data <> [] && circle_contains c p
      then List.fold_left (fun acc d -> f p d acc) accum data
      else accum
    | QNode (nw, ne, sw, se) ->
      let accum = fold_circ f accum nw c in
      let accum = fold_circ f accum ne c in
      let accum = fold_circ f accum sw c in
      fold_circ f accum se c

end