552 lines
23 KiB
Fortran
552 lines
23 KiB
Fortran
! ==============================================================================
|
|
! Module : kdtree_mod
|
|
! Purpose: k-d tree construction and nearest-neighbour search for k-dimensional
|
|
! point sets.
|
|
!
|
|
! A k-d tree is a binary space-partitioning tree. At depth d the splitting
|
|
! axis cycles as axis = mod(d, k) + 1. The median point along the chosen axis
|
|
! becomes the node; the smaller and larger halves recurse left and right.
|
|
! The result is a balanced tree of height ceil(log2 n).
|
|
!
|
|
! Nearest-neighbour search:
|
|
! Descend to the leaf that would contain the query, then unwind. At each
|
|
! node the algorithm checks whether the splitting hyperplane is close enough
|
|
! that the far subtree might contain a nearer point; if so, that subtree is
|
|
! also searched. This gives O(log n) average complexity while guaranteeing
|
|
! the true nearest neighbour.
|
|
!
|
|
! Node-visit counting:
|
|
! A node is counted as visited the moment its coordinates or split axis are
|
|
! first read (i.e., at the top of nn_search for that node index).
|
|
!
|
|
! Public interface:
|
|
! max_dim -- compile-time maximum dimensionality (= 3)
|
|
! kdnode -- derived type: coords(max_dim), left, right, axis
|
|
! pool(1:pool_n) -- flat array in which the tree lives
|
|
! pool_n -- number of allocated nodes (= number of points)
|
|
! pool_reset(n) -- allocate pool for n points; call before build_kdtree
|
|
! build_kdtree(pts, n_pts, idx, n, k) -> root node index
|
|
! nn_search(root, query, k, best_dist2, best_idx, visited)
|
|
! ==============================================================================
|
|
module kdtree_mod
|
|
implicit none
|
|
private
|
|
|
|
integer, parameter, public :: max_dim = 3 ! max k supported at compile time
|
|
|
|
! --------------------------------------------------------------------------
|
|
! One node in the k-d tree.
|
|
! coords -- coordinates of the point stored here (max_dim elements)
|
|
! left -- pool index of the left child (0 = no child)
|
|
! right -- pool index of the right child (0 = no child)
|
|
! axis -- splitting axis for this node (1-based; cycles 1..k by depth)
|
|
! --------------------------------------------------------------------------
|
|
type, public :: kdnode
|
|
real :: coords(max_dim) = 0.0
|
|
integer :: left = 0
|
|
integer :: right = 0
|
|
integer :: axis = 0
|
|
end type
|
|
|
|
! Pool is allocatable: sized exactly to the point count at build time.
|
|
! This replaces the old compile-time max_pool constant, which caused silent
|
|
! memory corruption when the static array was smaller than the point count.
|
|
type(kdnode), allocatable, public :: pool(:)
|
|
integer, public :: pool_n = 0
|
|
|
|
public :: pool_reset, build_kdtree, nn_search
|
|
|
|
contains
|
|
|
|
! --------------------------------------------------------------------------
|
|
! pool_reset(n) -- (re)allocate the pool for exactly n points and reset
|
|
! the node counter. Call once before each build_kdtree.
|
|
! --------------------------------------------------------------------------
|
|
subroutine pool_reset(n)
|
|
integer, intent(in) :: n
|
|
if (allocated(pool)) deallocate(pool)
|
|
allocate(pool(n))
|
|
pool_n = 0
|
|
end subroutine
|
|
|
|
! --------------------------------------------------------------------------
|
|
! alloc_node -- claim the next pool slot, initialise it, return its index.
|
|
! --------------------------------------------------------------------------
|
|
integer function alloc_node(coords, axis)
|
|
real, intent(in) :: coords(max_dim) ! full coordinate column
|
|
integer, intent(in) :: axis ! splitting axis for this node
|
|
pool_n = pool_n + 1
|
|
alloc_node = pool_n
|
|
pool(pool_n)%coords = coords
|
|
pool(pool_n)%axis = axis
|
|
pool(pool_n)%left = 0
|
|
pool(pool_n)%right = 0
|
|
end function
|
|
|
|
! --------------------------------------------------------------------------
|
|
! isort -- insertion-sort idx(1:n) ascending by pts(dim, idx(i)).
|
|
! O(n^2) but branchless and cache-friendly for tiny n; used only as the
|
|
! base case of qsort for partitions of size <= 20.
|
|
! --------------------------------------------------------------------------
|
|
subroutine isort(idx, n, pts, n_pts, dim)
|
|
integer, intent(inout) :: idx(n)
|
|
integer, intent(in) :: n, n_pts, dim
|
|
real, intent(in) :: pts(max_dim, n_pts)
|
|
integer :: i, j, tmp
|
|
do i = 2, n
|
|
tmp = idx(i)
|
|
j = i - 1
|
|
do while (j >= 1)
|
|
if(pts(dim, idx(j)) <= pts(dim, tmp))exit
|
|
idx(j+1) = idx(j)
|
|
j = j - 1
|
|
end do
|
|
idx(j+1) = tmp
|
|
end do
|
|
end subroutine
|
|
|
|
! --------------------------------------------------------------------------
|
|
! qsort -- iterative quicksort of idx(1:n) ascending by pts(dim, idx(i)).
|
|
!
|
|
! Non-recursive: uses an explicit integer stack (slo/shi) instead of the
|
|
! Fortran call stack. This is essential for large n because the recursive
|
|
! version caused stack overflow: build_kdtree (~19 deep) called recursive
|
|
! qsort (~19 deep), and gfortran allocated per-frame bookkeeping for the
|
|
! pts(max_dim, n_pts) dummy at each level, exhausting the Windows stack.
|
|
!
|
|
! Median-of-three pivot selection avoids O(n^2) on sorted input.
|
|
! Small partitions (<= 20 elements) are finished with isort.
|
|
!
|
|
! Stack depth is bounded to O(log n) by the "smaller-half-first" rule:
|
|
! push the LARGER partition first so the SMALLER sits on top and is
|
|
! processed next. For n = 400,000, log2(n) ~ 19; SMAX = 128 is ample.
|
|
! --------------------------------------------------------------------------
|
|
subroutine qsort(idx, n, pts, n_pts, dim)
|
|
integer, intent(inout) :: idx(n)
|
|
integer, intent(in) :: n, n_pts, dim
|
|
real, intent(in) :: pts(max_dim, n_pts)
|
|
|
|
integer, parameter :: SMAX = 128
|
|
integer :: slo(SMAX), shi(SMAX) ! explicit partition-boundary stack
|
|
integer :: top, lo, hi, p, i, mid, tmp
|
|
real :: pivot
|
|
|
|
if (n <= 1) return
|
|
|
|
top = 1
|
|
slo(top) = 1
|
|
shi(top) = n
|
|
|
|
do while (top > 0)
|
|
lo = slo(top)
|
|
hi = shi(top)
|
|
top = top - 1
|
|
|
|
! Base case: insertion sort for small partitions
|
|
if (hi - lo < 20) then
|
|
call isort(idx(lo:hi), hi - lo + 1, pts, n_pts, dim)
|
|
cycle
|
|
end if
|
|
|
|
! Median-of-three: compare lo, mid, hi; place median at hi as pivot
|
|
mid = lo + (hi - lo) / 2
|
|
if (pts(dim,idx(lo)) > pts(dim,idx(mid))) then; tmp=idx(lo); idx(lo)=idx(mid); idx(mid)=tmp; end if
|
|
if (pts(dim,idx(lo)) > pts(dim,idx(hi))) then; tmp=idx(lo); idx(lo)=idx(hi); idx(hi)=tmp; end if
|
|
if (pts(dim,idx(mid)) > pts(dim,idx(hi))) then; tmp=idx(mid); idx(mid)=idx(hi); idx(hi)=tmp; end if
|
|
tmp = idx(mid); idx(mid) = idx(hi); idx(hi) = tmp ! median -> pivot slot
|
|
|
|
! Lomuto partition on idx(lo:hi)
|
|
pivot = pts(dim, idx(hi))
|
|
p = lo - 1
|
|
do i = lo, hi - 1
|
|
if (pts(dim, idx(i)) <= pivot) then
|
|
p = p + 1
|
|
tmp = idx(p); idx(p) = idx(i); idx(i) = tmp
|
|
end if
|
|
end do
|
|
p = p + 1
|
|
tmp = idx(p); idx(p) = idx(hi); idx(hi) = tmp ! place pivot at p
|
|
|
|
! Push sub-problems: LARGER partition goes below SMALLER (O(log n) stack)
|
|
if ((p - lo) >= (hi - p)) then ! left [lo,p-1] is larger
|
|
if (lo < p - 1) then; top=top+1; slo(top)=lo; shi(top)=p-1; end if
|
|
if (p + 1 < hi) then; top=top+1; slo(top)=p+1; shi(top)=hi; end if
|
|
else ! right [p+1,hi] is larger
|
|
if (p + 1 < hi) then; top=top+1; slo(top)=p+1; shi(top)=hi; end if
|
|
if (lo < p - 1) then; top=top+1; slo(top)=lo; shi(top)=p-1; end if
|
|
end if
|
|
end do
|
|
end subroutine
|
|
|
|
! --------------------------------------------------------------------------
|
|
! build_kdtree -- iterative balanced k-d tree construction.
|
|
!
|
|
! pts(max_dim, n_pts) : full point array (not modified)
|
|
! idx(n) : working index array (rearranged in-place)
|
|
! n : number of points (0 = empty tree)
|
|
! k : dimensionality (1..max_dim)
|
|
!
|
|
! Returns the pool index of the root node (0 if n = 0).
|
|
!
|
|
! Algorithm -- iterative DFS using five parallel integer arrays as a stack.
|
|
! Each stack entry records the subproblem: index range [lo,hi] in idx,
|
|
! current depth (determines splitting axis), parent node index, and which
|
|
! child slot (0=left, 1=right) to write the new node index into.
|
|
!
|
|
! At each iteration:
|
|
! axis = mod(depth, k) + 1
|
|
! sort idx(lo:hi) along that axis
|
|
! median = lo + (hi-lo)/2 (lower median, 1-based in the full idx array)
|
|
! create a node for idx(median), wire it into the parent's child slot
|
|
! push left sub-problem [lo, median-1] if non-empty
|
|
! push right sub-problem [median+1, hi ] if non-empty
|
|
!
|
|
! Stack depth is O(log n); SMAX=64 handles up to 2^63 points.
|
|
! --------------------------------------------------------------------------
|
|
integer function build_kdtree(pts, n_pts, idx, n, k) result(root_nd)
|
|
integer, intent(in) :: n_pts, n, k
|
|
real, intent(in) :: pts(max_dim, n_pts)
|
|
integer, intent(inout) :: idx(n)
|
|
|
|
integer, parameter :: SMAX = 64
|
|
integer :: st_lo(SMAX), st_hi(SMAX), st_dep(SMAX)
|
|
integer :: st_par(SMAX), st_sid(SMAX) ! parent node index; child side
|
|
integer :: top, lo, hi, dep, par, sid
|
|
integer :: axis, med, nd
|
|
|
|
root_nd = 0
|
|
if (n == 0) return
|
|
|
|
top = 1
|
|
st_lo(top) = 1; st_hi(top) = n
|
|
st_dep(top)= 0; st_par(top)= 0; st_sid(top)= 0
|
|
|
|
do while (top > 0)
|
|
lo = st_lo(top); hi = st_hi(top)
|
|
dep = st_dep(top); par = st_par(top); sid = st_sid(top)
|
|
top = top - 1
|
|
|
|
axis = mod(dep, k) + 1
|
|
call qsort(idx(lo:hi), hi - lo + 1, pts, n_pts, axis)
|
|
med = lo + (hi - lo) / 2 ! median index in full idx array
|
|
|
|
nd = alloc_node(pts(:, idx(med)), axis)
|
|
|
|
! Wire this node into the tree
|
|
if (par == 0) then
|
|
root_nd = nd ! first node popped is the root
|
|
else if (sid == 0) then
|
|
pool(par)%left = nd
|
|
else
|
|
pool(par)%right = nd
|
|
end if
|
|
|
|
! Push non-empty child sub-problems (left first so right is on top
|
|
! and processed next -- keeps the near branch hot in cache)
|
|
if (lo < med) then
|
|
top = top + 1
|
|
st_lo(top) = lo; st_hi(top) = med - 1
|
|
st_dep(top)= dep+1; st_par(top)= nd; st_sid(top)= 0
|
|
end if
|
|
if (med < hi) then
|
|
top = top + 1
|
|
st_lo(top) = med+1; st_hi(top) = hi
|
|
st_dep(top)= dep+1; st_par(top)= nd; st_sid(top)= 1
|
|
end if
|
|
end do
|
|
end function
|
|
|
|
! --------------------------------------------------------------------------
|
|
! nn_search -- iterative nearest-neighbour search in a k-d tree.
|
|
!
|
|
! root : pool index of the tree root (0 = empty tree)
|
|
! query(k) : the point being searched for
|
|
! k : dimensionality
|
|
! best_dist2 : squared distance to the current best candidate (updated)
|
|
! best_idx : pool index of the current best candidate (updated)
|
|
! visited : running count of nodes accessed (updated)
|
|
!
|
|
! Algorithm -- iterative DFS with pruning, using two parallel stacks:
|
|
! stk_nd(i) : node index to visit
|
|
! stk_pd2(i) : squared distance from query to this subtree's splitting
|
|
! hyperplane (0.0 for the near subtree, split_diff^2 for far)
|
|
!
|
|
! At each iteration:
|
|
! 1. Pop (nd, plane_dist2). If plane_dist2 >= best_dist2, prune.
|
|
! 2. Visit nd: compute distance, update best.
|
|
! 3. Push far child with plane_dist2 = split_diff^2.
|
|
! Push near child with plane_dist2 = 0.0 (on top, visited first).
|
|
! The pruning check at pop time uses the CURRENT best_dist2, so it is
|
|
! at least as tight as the recursive check at call time.
|
|
!
|
|
! Stack depth O(log n); SMAX=128 is ample for any practical tree.
|
|
! --------------------------------------------------------------------------
|
|
subroutine nn_search(root, query, k, best_dist2, best_idx, visited)
|
|
integer, intent(in) :: root, k
|
|
real, intent(in) :: query(k)
|
|
real, intent(inout) :: best_dist2
|
|
integer, intent(inout) :: best_idx, visited
|
|
|
|
integer, parameter :: SMAX = 128
|
|
integer :: stk_nd(SMAX)
|
|
real :: stk_pd2(SMAX) ! lower-bound dist^2 for each entry
|
|
integer :: top, nd, near_child, far_child, i
|
|
real :: d2, split_diff
|
|
|
|
if (root == 0) return
|
|
|
|
top = 1
|
|
stk_nd(top) = root
|
|
stk_pd2(top) = 0.0
|
|
|
|
do while (top > 0)
|
|
nd = stk_nd(top)
|
|
top = top - 1
|
|
|
|
if (nd == 0) cycle
|
|
if (stk_pd2(top + 1) >= best_dist2) cycle ! prune: can't beat current best
|
|
|
|
! ---- Visit this node -----------------------------------------------
|
|
visited = visited + 1
|
|
|
|
d2 = 0.0
|
|
do i = 1, k
|
|
d2 = d2 + (query(i) - pool(nd)%coords(i))**2
|
|
end do
|
|
if (d2 < best_dist2) then
|
|
best_dist2 = d2
|
|
best_idx = nd
|
|
end if
|
|
|
|
! ---- Choose near / far children ------------------------------------
|
|
split_diff = query(pool(nd)%axis) - pool(nd)%coords(pool(nd)%axis)
|
|
if (split_diff <= 0.0) then
|
|
near_child = pool(nd)%left
|
|
far_child = pool(nd)%right
|
|
else
|
|
near_child = pool(nd)%right
|
|
far_child = pool(nd)%left
|
|
end if
|
|
|
|
! Push far child first (lower in stack), near child on top.
|
|
! Near child always has plane_dist2 = 0 (query is on that side).
|
|
! Far child has plane_dist2 = split_diff^2; pruned at pop time if
|
|
! best_dist2 has improved enough since this push.
|
|
if (far_child /= 0) then
|
|
top = top + 1
|
|
stk_nd(top) = far_child
|
|
stk_pd2(top) = split_diff**2
|
|
end if
|
|
if (near_child /= 0) then
|
|
top = top + 1
|
|
stk_nd(top) = near_child
|
|
stk_pd2(top) = 0.0
|
|
end if
|
|
end do
|
|
end subroutine
|
|
|
|
end module kdtree_mod
|
|
|
|
! ==============================================================================
|
|
! Program : kd_demo
|
|
! Purpose : Demonstrates k-d tree construction and nearest-neighbour search on
|
|
! two datasets.
|
|
!
|
|
! Test 1 -- Wikipedia 2-D example:
|
|
! Points : (2,3) (5,4) (9,6) (4,7) (8,1) (7,2)
|
|
! Query : (9, 2)
|
|
! Expected: nearest = (8,1), distance = sqrt(2) ≈ 1.4142
|
|
!
|
|
! Test 2 -- 1000 random 3-D points uniformly distributed in [0,1)^3:
|
|
! Query : a random point also drawn from [0,1)^3
|
|
! A brute-force linear scan verifies the k-d tree result.
|
|
!
|
|
! Test 3 -- 400000 random 3-D points (same setup as Test 2, larger scale):
|
|
! Timings for tree build, k-d tree search, and brute-force search are
|
|
! reported so the O(log n) vs O(n) scaling can be observed directly.
|
|
!
|
|
! Output fields for each case:
|
|
! Searched for -- the query point
|
|
! Nearest neighbor -- the closest point found in the tree
|
|
! Distance -- Euclidean distance from query to nearest neighbor
|
|
! Nodes visited -- number of tree nodes accessed during the search
|
|
! ==============================================================================
|
|
program kd_demo
|
|
use kdtree_mod
|
|
implicit none
|
|
|
|
integer, parameter :: N2 = 6 ! Wikipedia 2-D example: 6 points
|
|
integer, parameter :: N3 = 1000 ! random 3-D test: 1000 points
|
|
integer, parameter :: N4 = 400000 ! large random 3-D test: 400000 points
|
|
|
|
! pts3, pts4, and idx are allocatable so they land on the heap rather than
|
|
! the stack. Under -O3 gfortran can stack-allocate program-unit locals;
|
|
! pts4 alone is 4.8 MB and idx is 1.6 MB, which together exhaust the
|
|
! typical 8 MB Windows stack long before any computation begins.
|
|
real :: pts2(max_dim, N2) ! 2-D points (tiny -- stack is fine)
|
|
real, allocatable :: pts3(:,:) ! 3-D random points (1000)
|
|
real, allocatable :: pts4(:,:) ! 3-D random points (400000)
|
|
integer, allocatable :: idx(:) ! working index scratch for builder
|
|
|
|
integer :: root ! pool index of tree root
|
|
integer :: best_idx ! pool index of nearest neighbour found
|
|
integer :: visited ! nodes visited during NN search
|
|
integer :: i
|
|
integer :: t0, t1, t2, t3, rate ! system_clock variables for timing
|
|
|
|
real :: query(max_dim) ! search point
|
|
real :: best_dist2 ! squared distance to best candidate
|
|
real :: best_dist ! Euclidean distance (sqrt of above)
|
|
|
|
! Allocate all large arrays on the heap up front
|
|
allocate(pts3(max_dim, N3))
|
|
allocate(pts4(max_dim, N4))
|
|
allocate(idx(max(N2, N4)))
|
|
|
|
! Seed the RNG from the system clock so results differ each run
|
|
call set_rng_seed()
|
|
|
|
! ==========================================================================
|
|
! Test 1: Wikipedia 2-D example
|
|
! ==========================================================================
|
|
pts2 = 0.0 ! zero-fill; unused 3rd coord stays 0
|
|
pts2(1,1) = 2.0; pts2(2,1) = 3.0
|
|
pts2(1,2) = 5.0; pts2(2,2) = 4.0
|
|
pts2(1,3) = 9.0; pts2(2,3) = 6.0
|
|
pts2(1,4) = 4.0; pts2(2,4) = 7.0
|
|
pts2(1,5) = 8.0; pts2(2,5) = 1.0
|
|
pts2(1,6) = 7.0; pts2(2,6) = 2.0
|
|
|
|
call pool_reset(N2)
|
|
idx(1:N2) = [(i, i = 1, N2)]
|
|
root = build_kdtree(pts2, N2, idx(1:N2), N2, 2)
|
|
|
|
query = 0.0
|
|
query(1) = 9.0
|
|
query(2) = 2.0
|
|
|
|
best_dist2 = huge(best_dist2)
|
|
best_idx = 0
|
|
visited = 0
|
|
call nn_search(root, query(1:2), 2, best_dist2, best_idx, visited)
|
|
best_dist = sqrt(best_dist2)
|
|
|
|
write(*, '(a)') '--- Wikipedia 2-D Example ---'
|
|
write(*, '(a,2f6.1)') 'Searched for: ', query(1:2)
|
|
write(*, '(a,2f6.1)') 'Nearest neighbor: ', pool(best_idx)%coords(1:2)
|
|
write(*, '(a,f9.4)') 'Distance: ', best_dist
|
|
write(*, '(a,i0)') 'Nodes visited: ', visited
|
|
|
|
! ==========================================================================
|
|
! Test 2: 1000 random 3-D points in the unit cube [0,1)^3
|
|
! ==========================================================================
|
|
call random_number(pts3) ! all 3 * 1000 values filled uniformly
|
|
|
|
call pool_reset(N3)
|
|
idx(1:N3) = [(i, i = 1, N3)]
|
|
root = build_kdtree(pts3, N3, idx(1:N3), N3, 3)
|
|
|
|
call random_number(query(1:3)) ! query point also in [0,1)^3
|
|
|
|
best_dist2 = huge(best_dist2)
|
|
best_idx = 0
|
|
visited = 0
|
|
call nn_search(root, query(1:3), 3, best_dist2, best_idx, visited)
|
|
best_dist = sqrt(best_dist2)
|
|
|
|
write(*, '(a)') ''
|
|
write(*, '(a)') '--- 1000 Random 3-D Points (unit cube) ---'
|
|
write(*, '(a,3f9.5)') 'Searched for: ', query(1:3)
|
|
write(*, '(a,3f9.5)') 'Nearest neighbor: ', pool(best_idx)%coords(1:3)
|
|
write(*, '(a,f9.6)') 'Distance: ', best_dist
|
|
write(*, '(a,i0)') 'Nodes visited: ', visited
|
|
|
|
! Brute-force linear scan to verify correctness
|
|
call brute_force_nn(pts3, N3, query(1:3), 3)
|
|
|
|
! ==========================================================================
|
|
! Test 3: 400000 random 3-D points in the unit cube [0,1)^3, with timing.
|
|
! Three phases are timed separately:
|
|
! build -- construct the k-d tree from 400000 points
|
|
! search -- single nearest-neighbour query via k-d tree
|
|
! brute -- exhaustive O(n) scan for the same query (reference)
|
|
! ==========================================================================
|
|
call random_number(pts4) ! 3 * 400000 uniform values in [0,1)
|
|
|
|
! -- Build -----------------------------------------------------------------
|
|
call pool_reset(N4)
|
|
idx(1:N4) = [(i, i = 1, N4)]
|
|
call system_clock(t0, rate)
|
|
root = build_kdtree(pts4, N4, idx(1:N4), N4, 3)
|
|
call system_clock(t1)
|
|
|
|
! -- k-d tree search -------------------------------------------------------
|
|
call random_number(query(1:3))
|
|
|
|
best_dist2 = huge(best_dist2)
|
|
best_idx = 0
|
|
visited = 0
|
|
call system_clock(t2)
|
|
call nn_search(root, query(1:3), 3, best_dist2, best_idx, visited)
|
|
call system_clock(t3)
|
|
best_dist = sqrt(best_dist2)
|
|
|
|
write(*, '(a)') ''
|
|
write(*, '(a)') '--- 400000 Random 3-D Points (unit cube) ---'
|
|
write(*, '(a,f9.4,a)') 'Build time: ', real(t1-t0)/real(rate)*1000.0, ' ms'
|
|
write(*, '(a,3f9.5)') 'Searched for: ', query(1:3)
|
|
write(*, '(a,3f9.5)') 'Nearest neighbor: ', pool(best_idx)%coords(1:3)
|
|
write(*, '(a,f9.6)') 'Distance: ', best_dist
|
|
write(*, '(a,i0)') 'Nodes visited: ', visited
|
|
write(*, '(a,f9.4,a)') 'Search time (k-d): ', real(t3-t2)/real(rate)*1e6, ' us'
|
|
|
|
! -- Brute-force (timed) ---------------------------------------------------
|
|
call system_clock(t2)
|
|
call brute_force_nn(pts4, N4, query(1:3), 3)
|
|
call system_clock(t3)
|
|
write(*, '(a,f12.4,a)') 'Search time (BF): ', real(t3-t2)/real(rate)*1e6, ' us'
|
|
|
|
contains
|
|
|
|
! --------------------------------------------------------------------------
|
|
! set_rng_seed -- seed Fortran's intrinsic RNG from the system clock.
|
|
! Each element of the seed array is shifted by 37*(i-1) so adjacent
|
|
! elements differ, reducing the chance of a degenerate seed state.
|
|
! --------------------------------------------------------------------------
|
|
subroutine set_rng_seed()
|
|
integer :: clock, sz, j
|
|
integer, allocatable :: seed(:)
|
|
call system_clock(clock)
|
|
call random_seed(size=sz)
|
|
allocate(seed(sz))
|
|
seed = clock + 37 * [(j - 1, j = 1, sz)]
|
|
call random_seed(put=seed)
|
|
deallocate(seed)
|
|
end subroutine
|
|
|
|
! --------------------------------------------------------------------------
|
|
! brute_force_nn -- O(n) exhaustive nearest-neighbour scan.
|
|
! Used to verify that the k-d tree search returned the correct answer.
|
|
! --------------------------------------------------------------------------
|
|
subroutine brute_force_nn(pts, n, query, k)
|
|
integer, intent(in) :: n, k
|
|
real, intent(in) :: pts(max_dim, n), query(k)
|
|
integer :: i, j, best
|
|
real :: d2, best_d2
|
|
best_d2 = huge(best_d2)
|
|
best = 0
|
|
do i = 1, n
|
|
d2 = 0.0
|
|
do j = 1, k
|
|
d2 = d2 + (query(j) - pts(j, i))**2
|
|
end do
|
|
if (d2 < best_d2) then
|
|
best_d2 = d2
|
|
best = i
|
|
end if
|
|
end do
|
|
write(*, '(a,3f9.5)') 'Brute-force check: ', pts(1:3, best)
|
|
write(*, '(a,f9.6)') 'Brute-force dist: ', sqrt(best_d2)
|
|
end subroutine
|
|
|
|
end program kd_demo
|