Simple Static KD-Trees using Julia

是 否 +2 论坛币

k人 参与回答

经管之家送您一份

应届毕业生专属福利!

求职就业群

赵安豆老师微信：zhaoandou666

经管之家联合CDA

送您一个全额奖学金名额~ !

立即领取

感谢您参与论坛问题回答

经管之家送您两个论坛币！

+2 论坛币

1. # Simple static kd-trees.
   
2. 
   
3. 
   
4. abstract KDNode
   
5. 
   
6. 
   
7. immutable KDTree{T <: FloatingPoint}
   
8.         # A matrix of n, m-dimensional observations
   
9.         xs::AbstractMatrix{T}
   
10. 
    
11.         # Permutaiton of the data used by the structure (to avoid sorting or
    
12.         # otherwise modifying xs)
    
13.         perm::Vector{Int}
    
14. 
    
15.         root::KDNode
    
16. 
    
17.         verts::Set{Vector{T}}
    
18. 
    
19.         # Top-level bounding box
    
20.         bounds::Matrix{T}
    
21. end
    
22. 
    
23. 
    
24. # Construct a kd-tree
    
25. #
    
26. # Args:
    
27. #   xs: Data organized in the tree.
    
28. #   leaf_size_factor: Stop spliting if a node contains
    
29. #       fewer than leaf_size_factor * n elements.
    
30. #   leaf_diameter_factor: Stop spliting if a node's bounding
    
31. #       hypercube has diameter less that leaf_diameter_factor
    
32. #       times the diameter of the root node's bounding hypercube.
    
33. #
    
34. # Returns:
    
35. #   A KDTree object
    
36. #
    
37. function KDTree{T <: FloatingPoint}(xs::AbstractMatrix{T},
    
38.                                         leaf_size_factor=0.05,
    
39.                                         leaf_diameter_factor=0.0)
    
40.         n, m = size(xs)
    
41.         perm = collect(1:n)
    
42. 
    
43.         bounds = Array(T, (2, m))
    
44.         for j in 1:m
    
45.                 col = xs[:,j]
    
46.                 bounds[1, j] = minimum(col)
    
47.                 bounds[2, j] = maximum(col)
    
48.         end
    
49.         diam = diameter(bounds)
    
50. 
    
51.         leaf_size_cutoff = iceil(leaf_size_factor * n)
    
52.         leaf_diameter_cutoff = leaf_diameter_factor * diam
    
53.         verts = Set{Vector{T}}()
    
54. 
    
55.         # Add verticies defined by the bounds
    
56.         for vert in product([bounds[:,j] for j in 1:m]...)
    
57.                 push!(verts, T[vert...])
    
58.         end
    
59. 
    
60.         root = build_kdtree(xs, sub(perm, 1:n), bounds,
    
61.                                 leaf_size_cutoff, leaf_diameter_cutoff, verts)
    
62. 
    
63.         KDTree(xs, collect(1:n), root, verts, bounds)
    
64. end
    
65. 
    
66. 
    
67. immutable KDLeafNode <: KDNode
    
68. end
    
69. 
    
70. 
    
71. immutable KDInternalNode{T <: FloatingPoint} <: KDNode
    
72.         j::Int # dimension on which the data is split
    
73.         med::T # median value where the split occours
    
74.         leftnode::KDNode
    
75.         rightnode::KDNode
    
76. end
    
77. 
    
78. 
    
79. 
    
80. # Compute the diamater of a hypercube defined by bounds
    
81. #
    
82. # Args:
    
83. #   bounds: A 2 by n matrix where bounds[1,i] gives the lower bound in
    
84. #           the ith dimension and bonuds[2,j] the upper.
    
85. #
    
86. # Returns:
    
87. #   Computed diameter
    
88. #
    
89. function diameter(bounds::Matrix)
    
90.         euclidean(vec(bounds[1,:]), vec(bounds[2,:]))
    
91. end
    
92. 
    
93. 
    
94. # Recursively build a kd-tree
    
95. #
    
96. # Args:
    
97. #   xs: Data being orginized.
    
98. #   perm: Permutation of the data, used to avoid
    
99. #       directly sorting or modifying xs.
    
100. #   bounds: Bounding hypercube of the node.
     
101. #   leaf_size_cutoff: stop spliting on nodes with more than
     
102. #       this many values.
     
103. #   leaf_diameter_cutoff: stop splitting on nodes with less
     
104. #        than this diameter.
     
105. #   verts: current set of vertexes
     
106. #
     
107. # Modifies:
     
108. #   perm, verts
     
109. #
     
110. # Returns:
     
111. #   Either a KDLeafNode or a KDInternalNode
     
112. #
     
113. function build_kdtree{T}(xs::AbstractMatrix{T},
     
114.                              perm::SubArray,
     
115.                              bounds::Matrix{T},
     
116.                              leaf_size_cutoff::Int,
     
117.                              leaf_diameter_cutoff::T,
     
118.                              verts::Set{Vector{T}})
     
119.         n, m = size(xs)
     
120. 
     
121.         if length(perm) <= leaf_size_cutoff || diameter(bounds) <= leaf_diameter_cutoff
     
122.                 return KDLeafNode()
     
123.         end
     
124. 
     
125.         # split on the dimension with the largest spread
     
126.         j = 1
     
127.         maxspread = 0
     
128.         for k in 1:m
     
129.                 xmin = Inf
     
130.                 xmax = -Inf
     
131.                 for i in perm
     
132.                         xmin = min(xmin, xs[i, k])
     
133.                         xmax = max(xmax, xs[i, k])
     
134.                 end
     
135.                 if xmax - xmin > maxspread
     
136.                         maxspread = xmax - xmin
     
137.                         j = k
     
138.                 end
     
139.         end
     
140. 
     
141.         # find the median and partition
     
142.         if isodd(length(perm))
     
143.                 mid = div(length(perm), 2)
     
144.                 select!(perm, mid, by=i -> xs[i, j])
     
145.                 med = xs[perm[mid], j]
     
146.                 mid1 = mid
     
147.                 mid2 = mid + 1
     
148.         else
     
149.                 mid1 = div(length(perm), 2)
     
150.                 mid2 = mid1 + 1
     
151.                 select!(perm, mid1, by=i -> xs[i, j])
     
152.                 select!(perm, mid2, by=i -> xs[i, j])
     
153.                 med = (xs[perm[mid1], j] + xs[perm[mid2], j]) / 2
     
154.         end
     
155. 
     
156.         leftbounds = copy(bounds)
     
157.         leftbounds[2, j] = med
     
158.         leftnode = build_kdtree(xs, sub(perm, 1:mid1), bounds,
     
159.                                     leaf_size_cutoff, leaf_diameter_cutoff, verts)
     
160. 
     
161.         rightbounds = copy(bounds)
     
162.         rightbounds[1, j] = med
     
163.         rightnode = build_kdtree(xs, sub(perm, mid2:length(perm)), bounds,
     
164.                                      leaf_size_cutoff, leaf_diameter_cutoff, verts)
     
165. 
     
166.         coords = Array(Array, m)
     
167.         for i in 1:m
     
168.                 if i == j
     
169.                         coords[i] = [med]
     
170.                 else
     
171.                         coords[i] = bounds[:, i]
     
172.                 end
     
173.         end
     
174. 
     
175.         for vert in product(coords...)
     
176.                 push!(verts, T[vert...])
     
177.         end
     
178. 
     
179.         KDInternalNode{T}(j, med, leftnode, rightnode)
     
180. end
     
181. 
     
182. 
     
183. # Given a bounding hypecube, return its verticies
     
184. function bounds_verts(bounds::Matrix)
     
185.         collect(product([bounds[:, i] for i in 1:size(bounds, 2)]...))
     
186. end
     
187. 
     
188. 
     
189. # Traverse the tree to the bottom and return the verticies of
     
190. # the leaf node's bounding hypercube.
     
191. function traverse{T}(kdtree::KDTree{T}, xs::AbstractVector{T})
     
192.         m = size(kdtree.bounds, 2)
     
193. 
     
194.         if length(xs) != m
     
195.                 error("$(m)-dimensional kd-tree searched with a length $(length(x)) vector.")
     
196.         end
     
197. 
     
198.         for j in 1:m
     
199.                 if xs[j] < kdtree.bounds[1, j] || xs[j] > kdtree.bounds[2, j]
     
200.                         error(
     
201.                                 """
     
202.                                 Loess cannot perform extrapolation. Predict can only be applied
     
203.                                 to points within the bounding hypercube of the data used to train
     
204.                                 the model.
     
205.                                 """)
     
206.                 end
     
207.         end
     
208. 
     
209.         bounds = copy(kdtree.bounds)
     
210.         node = kdtree.root
     
211.         while !isa(node, KDLeafNode)
     
212.                 if xs[node.j] <= node.med
     
213.                         bounds[2, node.j] = node.med
     
214.                         node = node.leftnode
     
215.                 else
     
216.                         bounds[1, node.j] = node.med
     
217.                         node = node.rightnode
     
218.                 end
     
219.         end
     
220. 
     
221.         bounds_verts(bounds)
     
222. end

复制代码

扫码加我 拉你入群

请注明：姓名-公司-职位

以便审核进群资格，未注明则拒绝

分享0 收藏0 回帖

关键词：static simple Trees Julia Using