1%% 2%% e3d_kd3.erl -- 3%% 4%% Kd tree implementation. 5%% 6%% Copyright (c) 2009-2011 Dan Gudmundsson 7%% 8%% See the file "license.terms" for information on usage and redistribution 9%% of this file, and for a DISCLAIMER OF ALL WARRANTIES. 10%% 11 12%%% @doc kd-tree from wikipedia 1 Jul 2009 13%%% Builds and traverses kd-tree of dimension 3 i.e. keys are {X,Y,Z} 14%%% 15%%% Each object is a 2-tuple, {point(), value} where point is a tuple 16%%% of size 3 and value is any term. 17%%% 18%%% Currently without any balencing or optimizations 19%%% 20%%% @end 21 22%% Implementation: 23%% I have placed the data in the leafs and build a slightly larger 24%% record in the nodes to avoid re-balancing nodes after removing 25%% an element which is needed standard algorithm. 26%% Since take_nearest is probably the most used function in the api. 27 28-module(e3d_kd3). 29 30-export([from_list/1, to_list/1, 31 empty/0,is_empty/1,is_kd3/1,size/1, 32 enter/3, update/3, 33 delete/2, delete_object/2, 34 nearest/2, take_nearest/2, 35 fold/5 36%% map/2 37 ]). 38 39-include("e3d.hrl"). %% For types 40 41-define(REBALANCE, 0.7). %% Rebalance after % deletions 42-define(NODE(Point,Axis,Left,Right), {Point,Axis,Left,Right}). 43 44-type kd3_object() :: {term(), e3d_vec:point()}. 45 46%% Internal 47-type tree() :: [kd3_object()] | 48 {Med :: float(), 49 Axis :: integer(), 50 Left :: tree(), 51 Right :: tree()}. 52 53-record(e3d_kd3, { tree=[] :: tree() }). 54-type e3d_kd3() :: #e3d_kd3{}. 55 56-export_type([e3d_kd3/0]). 57 58%%% @doc Returns an empty Tree. 59-spec empty() -> e3d_kd3(). 60empty() -> #e3d_kd3{tree=[]}. 61 62%%% @doc Returns true if Tree is an empty tree. 63-spec is_empty(e3d_kd3()) -> boolean(). 64is_empty(#e3d_kd3{tree=[]}) -> true; 65is_empty(#e3d_kd3{}) -> false. 66 67-spec is_kd3(e3d_kd3()) -> boolean(). 68is_kd3(#e3d_kd3{}) -> true; 69is_kd3(_) -> false. 70 71%%% @doc Returns the number of objects in the Tree. 72%%% Note: This function currently traverses the tree. 73-spec size(e3d_kd3()) -> integer(). 74size(#e3d_kd3{tree=Tree}) -> 75 size_1(Tree, 0). 76 77size_1({_, _, L0, R0}, Sz) -> 78 size_1(L0, size_1(R0, Sz)); 79size_1(List, Sz) -> 80 Sz + length(List). 81 82%%% @doc Removes all the node(s) with key Point from Tree1 and returns the new Tree2 83%%% assumes the Key is present otherwise it crashes. 84-spec delete(e3d_vec:point(), e3d_kd3()) -> e3d_kd3(). 85delete({_,_,_} = Point, #e3d_kd3{tree=Tree}) -> 86 #e3d_kd3{tree=delete_1(Point, Tree)}. 87 88delete_1(Key,{Med, Axis, L0, R0}) -> 89 case Med < element(Axis, Key) of 90 true -> 91 R = delete_1(Key, R0), 92 delete_2(Med, Axis, L0, R); 93 false -> 94 L = delete_1(Key, L0), 95 delete_2(Med, Axis, L, R0) 96 end; 97delete_1(Key, [{_,Key}|_]) -> []. 98 99delete_2(_, _, [], R) -> R; 100delete_2(_, _, L, []) -> L; 101delete_2(Med, Axis, L, R) -> ?NODE(Med,Axis,L,R). 102 103%%% @doc Removes the Object from Tree1 and returns the new Tree2 104%%% crashes if object is not present. 105-spec delete_object(kd3_object(), e3d_kd3()) -> e3d_kd3(). 106delete_object({_,{_,_,_}} = Object, _T=#e3d_kd3{tree=Tree}) -> 107 #e3d_kd3{tree=delete_object_1(Object, Tree)}. 108 109delete_object_1({_,Key}=Object,{Med, Axis, L0, R0}) -> 110 case Med < element(Axis, Key) of 111 true -> 112 R = delete_object_1(Object, R0), 113 delete_object_2(Med, Axis, L0, R); 114 false -> 115 L = delete_object_1(Object, L0), 116 delete_object_2(Med, Axis, L, R0) 117 end; 118delete_object_1(Object, Leaf) -> 119 delete_object_3(Object, Leaf, []). 120 121delete_object_2(_, _, [], R) -> R; 122delete_object_2(_, _, L, []) -> L; 123delete_object_2(Med, Axis, L, R) -> ?NODE(Med,Axis,L,R). 124 125delete_object_3(Object, [Object|R], Acc) -> Acc ++ R; 126delete_object_3(Object, [H|T], Acc) -> 127 delete_object_3(Object, T, [H|Acc]). 128 129%% rebalence(Tree = #e3d_kd3{deleted=D, total=T}) when T > 1000, D/T > ?REBALANCE -> 130%% from_list(to_list(Tree)); 131%%rebalence(T) -> T. 132 133%%% @doc Builds a kd-tree from a list of objects. 134-spec from_list([kd3_object()]) -> e3d_kd3(). 135from_list([]) -> 136 #e3d_kd3{tree=[]}; 137from_list(List) -> 138 {_N,BB} = lists:foldl(fun({_, Point}, {N,BB}) -> {N+1,e3d_bv:union(BB, Point)} end, 139 {0,e3d_bv:box()}, List), 140 #e3d_kd3{tree=from_list(List, BB)}. %, total=N}. 141 142from_list(List = [_|_], BB) -> 143 case e3d_bv:max_extent(BB) of 144 undefined -> %% All positions are exactly the same 145 List; 146 Axis -> 147 Split = element(Axis, e3d_bv:center(BB)), 148 {Left,LBB,Right,RBB} = split(List, Axis, Split), 149 ?NODE(Split, Axis, %(Len bsl 2) bor Axis, 150 from_list(Left, LBB), 151 from_list(Right, RBB)) 152 end; 153from_list(Data, _) -> Data. 154 155%%% @doc Add node to the tree. 156-spec enter(e3d_vec:point(), term(), e3d_kd3()) -> e3d_kd3(). 157enter(Point, Data, #e3d_kd3{tree=Tree0}) -> 158 Tree = add_1({Data, Point}, Tree0), 159 #e3d_kd3{tree=Tree}. 160 161%%% @doc Updates the Object with Term from Tree1 and returns the new Tree2 162%%% crashes if object is not present. 163-spec update(kd3_object(), term(), e3d_kd3()) -> e3d_kd3(). 164update({_,{_,_,_}} = Object, Val, _T=#e3d_kd3{tree=Tree}) -> 165 #e3d_kd3{tree=update_1(Object, Val, Tree)}. 166 167update_1({_,Key}=Object,Val,{Med, Axis, L0, R0}) -> 168 case Med < element(Axis, Key) of 169 true -> 170 R = update_1(Object, Val, R0), 171 delete_object_2(Med, Axis, L0, R); 172 false -> 173 L = update_1(Object, Val, L0), 174 delete_object_2(Med, Axis, L, R0) 175 end; 176update_1(Object, Val, Leaf) -> 177 update_3(Object, Leaf, Val, []). 178 179update_3({_,Pos}=Object, [Object|R], Val, Acc) -> Acc ++ [{Val,Pos}|R]; 180update_3(Object, [H|T], Val, Acc) -> 181 update_3(Object, T, Val, [H|Acc]). 182 183%%% @doc Return all nodes in the tree. 184-spec to_list(e3d_kd3()) -> [kd3_object()]. 185to_list(#e3d_kd3{tree=Tree}) -> 186 to_list(Tree, []). 187 188to_list({_,_, L,R}, Acc0) -> 189 to_list(L,to_list(R,Acc0)); 190to_list([], Acc) -> Acc; 191to_list(Data, Acc) -> Data++Acc. 192 193%%% @spec (Point::tuple(), Tree::e3d_kd3()) -> {kd3_object(), e3d_kd3()} | undefined 194%%% @doc Returns the Object nearest the Point and a Tree with the Object deleted, 195%%% or undefined 196-spec take_nearest(e3d_vec:point(), Tree::e3d_kd3()) -> {kd3_object(), e3d_kd3()} | undefined. 197take_nearest({_,_,_}, #e3d_kd3{tree=[]}) -> undefined; 198take_nearest({_,_,_} = Point, Orig = #e3d_kd3{tree=Tree}) -> 199 [_|[Node|_]=_Closest] = nearest_1(Point, Tree, [undefined|undefined]), 200 %%Foo = [P1|| {P1,_} <- _Closest], 201 %io:format("Pick: ~p ~p~n=> ~p~n", [Point, Point =:= hd(Foo), Closest]), 202 {Node, delete_object(Node,Orig)}. 203 204%%% @doc Returns the object nearest the Key, or undefined if table is empty. 205-spec nearest(Key::e3d_vec:point(), Tree::e3d_kd3()) -> kd3_object() | undefined. 206nearest({_,_,_} = Point, #e3d_kd3{tree=Tree}) -> 207 [_|[Node|_]] = nearest_1(Point, Tree, [undefined|undefined]), 208 Node. 209 210nearest_1(Point, {SplitPos,Axis,L,R}, Closest) -> 211 PointPos = element(Axis, Point), 212 BorderDist = (PointPos - SplitPos), 213 Border = BorderDist * BorderDist, 214 case SplitPos < PointPos of 215 true -> nearest_2(Point, R, L, Border, Closest); 216 false -> nearest_2(Point, L, R, Border, Closest) 217 end; 218nearest_1(Point,[{_,Pos}|_] = Leaf, Closest0) -> 219 closest([e3d_vec:dist_sqr(Pos, Point)|Leaf],Closest0); 220nearest_1(_, [], Close) -> Close. 221 222nearest_2(Point, ThisSide, OtherSide, Border, Closest0) -> 223 case nearest_1(Point, ThisSide, Closest0) of 224 Closest = [Dist|_] when Dist < Border -> 225 Closest; 226 Closest -> 227 nearest_1(Point, OtherSide, Closest) 228 end. 229 230%%% @doc Traverse all Objects within Dist distance from point, notice that the order 231%%% is not specified. 232-spec fold(fun((kd3_object(),Acc::term()) -> term()), term(), 233 e3d_vec:point(), float(), e3d_kd3()) -> term(). 234fold(Fun, Acc, {_,_,_} = Point, Dist, #e3d_kd3{tree=Tree}) -> 235 fold_1(Tree, Point, Dist*Dist, Fun, Acc). 236 237fold_1({SplitPos,Axis,L,R}, Point, Dist2, Fun, Acc0) -> 238 PointPos = element(Axis, Point), 239 BorderDist = (PointPos - SplitPos), 240 Border = BorderDist * BorderDist, 241 if Border < Dist2 -> 242 Acc = fold_1(R, Point, Dist2, Fun, Acc0), 243 fold_1(L, Point, Dist2, Fun, Acc); 244 SplitPos < PointPos -> 245 fold_1(R, Point, Dist2, Fun, Acc0); 246 true -> 247 fold_1(L, Point, Dist2, Fun, Acc0) 248 end; 249fold_1([], _Point, _Dist2, _Fun, Acc) -> Acc; 250fold_1(List = [{_,Pos}|_], Point, Dist2, Fun, Acc) -> 251 case e3d_vec:dist_sqr(Pos, Point) =< Dist2 of 252 true -> 253 lists:foldl(Fun, Acc, List); 254 false -> 255 Acc 256 end. 257 258%%% Internal stuff %%% 259split(List, Axis, Pos) -> 260 split(List, Axis, Pos, [], e3d_bv:box(), [], e3d_bv:box()). 261 262split([H = {_,V}|T], Axis, Pos, L, LB, R, RB) -> 263 case element(Axis, V) =< Pos of 264 true -> 265 split(T, Axis, Pos, [H|L], e3d_bv:union(LB,V), R, RB); 266 false -> 267 split(T, Axis, Pos, L, LB, [H|R], e3d_bv:union(RB,V)) 268 end; 269split([], _, _, L,LB,R,RB) -> 270 {L,LB,R,RB}. 271 272closest([Dist1|_]=Node1, [Dist2|_]) when Dist1 < Dist2 -> Node1; 273closest(_, Node) -> Node. 274 275add_1(P0, []) -> [P0]; 276add_1({_,P0}=Obj, [{_,P1}|_]=T0) -> 277 BB = e3d_bv:box(P0, P1), 278 case e3d_bv:max_extent(BB) of 279 undefined -> %% All positions are exactly the same 280 [Obj|T0]; 281 Axis -> 282 Split = element(Axis, e3d_bv:center(BB)), 283 {Left,LBB,Right,RBB} = split([Obj|T0], Axis, Split), 284 ?NODE(Split, Axis, %(Len bsl 2) bor Axis, 285 from_list(Left, LBB), 286 from_list(Right, RBB)) 287 end; 288add_1({_,P0}=Obj, {SplitPos, Axis, L, R}) -> 289 PointPos = element(Axis, P0), 290 case SplitPos < PointPos of 291 true -> {SplitPos, Axis, L, add_1(Obj, R)}; 292 false -> {SplitPos, Axis, add_1(Obj, L), R} 293 end. 294 295 296