quad_tree.erl

%%%-------------------------------------------------------------------
%%% File    : quad_tree.erl
%%% @author Jamie Burrell
%%% @doc Erlang implementation of a quadratic tree.
%%% @end
%%% @since 2009-02-10
%%% @end
%%%-------------------------------------------------------------------
-module(quad_tree).

-export([new/1,add/2,remove/2,flatten/1,query_rect/2]).

-include("tree_records.hrl").

%% @spec new(Rect::rectangle()) -> quad_tree()
%% @type quad_tree() = #qnode{ bounds=rectangle(), point=point(), qnodes=[quad_tree()] }
%% @type rectangle() = #rectangle{ topleft=point(), bottomright=point() }
%% @type point() = #point{ x=number(), y=number() }
%% @doc Creates an empty quadratic tree, with the specified bounding rectangle.
new(Rect) ->
  #qnode{qnodes=nil, point=nil, bounds=Rect}.

%% @spec add(QTree::quad_tree(), Point::point()) -> quad_tree()
%% @doc Adds a point to a tree, if the tree bounds contain the point.
add(#qnode{bounds=Bounds} = QTree, Point) when is_record(Point, point) ->
  case rectangle:contains(Bounds, Point) of
    false -> QTree;
    _ -> add_valid(QTree, Point)
  end;
add([], Point) when is_record(Point, point) ->
  [];
add([H|T], Point) when is_record(Point, point) ->
  [add(H, Point)] ++ add(T, Point);
add(QTree, _Point) ->
  QTree.

%% add a point to the tree, when we know it is within the tree's bounds
% the case when we are adding a point to a leaf node (ie it has a single point already)
add_valid(#qnode{bounds=Bounds, qnodes=QNodes} = _QTree, Point) when is_list(QNodes) ->
  SubNodes = add(QNodes, Point),
#qnode{bounds=Bounds, point=nil, qnodes=SubNodes};
% the case when we are adding a point to a leaf node (ie it has a single point already)
add_valid(#qnode{bounds=Bounds, point=ExistingPoint} = QTree, #point{x=NewX,y=NewY} = Point) when is_record(ExistingPoint, point) ->
  #point{x=X,y=Y} = ExistingPoint,
  case {X, Y}  of
    {NewX, NewY} -> QTree;
    _ ->
      TmpSubNodes = add(branch(QTree), ExistingPoint),
      SubNodes = add(TmpSubNodes, Point),
      #qnode{bounds=Bounds, point=nil, qnodes=SubNodes}
  end;
% the case when we are adding a point to an empty parent node (ie it has no single point, and
% no subnodes)
add_valid(#qnode{bounds=Bounds, qnodes=QNodes} = _QTree, Point) ->
  #qnode{bounds=Bounds, point=Point, qnodes=QNodes}.

%% @spec remove(QTree::quad_tree(), Point::point()) -> quad_tree()
%% @doc Removes a point from the tree, if it exists.
remove(#qnode{bounds=Bounds} = QTree, Point) ->
  case rectangle:contains(Bounds, Point) of
    false -> QTree;
    _ -> remove_valid(QTree, Point)
  end;
remove([], _Point) ->
  [];
remove([H|T], Point) ->
  [remove(H, Point)] ++ remove(T, Point).

%% remove the point from the tree, when we know it is within the tree's bounds
remove_valid(#qnode{bounds=Bounds, qnodes=QNodes} = _QTree, Point) when is_list(QNodes) ->
  NewNodes = remove(QNodes, Point),
  flatten(#qnode{bounds=Bounds, qnodes=NewNodes, point=nil});
remove_valid(#qnode{bounds=Bounds, qnodes=QNodes, point=ExistingPoint} = QTree, #point{x=NewX, y=NewY} = _Point) when is_record(ExistingPoint,point) ->
  #point{x=X,y=Y} = ExistingPoint,
  case {X, Y}  of
    {NewX, NewY} ->
      #qnode{bounds=Bounds, qnodes=QNodes, point=nil};
    _ -> 
      QTree
  end;
remove_valid(QTree, _Point) ->
  QTree.

%% @spec query_rect(QTree::quad_tree(), Rect::rectangle()) -> [point()]
%% @doc Query for all points in the quad tree within a given rectangle.
query_rect(#qnode{bounds=Bounds} = QTree, Rect) ->
  case rectangle:intersects(Bounds, Rect) of
    false ->
      [];
    _ ->
      query_valid(QTree, Rect)
  end.

%% query this node for points, now we now it intersects the query rectangle
query_valid(#qnode{qnodes=QNodes} = _QTree, Rect) when is_list(QNodes) ->
  lists:flatten([ query_rect(QNode, Rect) || QNode <- QNodes ]); query_valid(#qnode{point=Point} = _QTree, Rect) when is_record(Point,point) ->
  case rectangle:contains(Rect, Point) of
    false ->
      [];
    _ ->
      [Point]
  end;
query_valid(_, _) ->
  [].

%% @spec flatten(QTree::quad_tree()) -> quad_tree()
%% @doc Flattens a tree whose subnodes together contain only a single point into a leaf
flatten(#qnode{bounds=Bounds, qnodes=QNodes} = QTree) when is_list(QNodes) ->
  TmpNodes = [flatten(Node) || Node <- QNodes],   AllPoints = query_rect(QTree, Bounds),   {NewNodes, NewPoint} =    case length(AllPoints) of      0 ->
  {nil, nil};
1 ->
  [OnlyPoint] = AllPoints,
  {nil, OnlyPoint};
_ ->
  {TmpNodes, nil}
end,
  #qnode{bounds=Bounds, qnodes=NewNodes, point=NewPoint};
flatten([]) ->
  [];
flatten([H|T]) ->
  [flatten(H)] ++ flatten(T);
flatten(QTree) ->
  QTree.

% split an existing node to get new subnodes
branch(#qnode{bounds=Bounds} = _QTree) ->
[new(Rect) || Rect <- rectangle:split(Bounds)].