aparler/RosettaCodeData
1
0
Fork 0
Code Issues 12 Pull requests 3 Projects Releases Packages Wiki Activity Actions

Data update

Browse source
This commit is contained in:
Ingy döt Net 2026-04-30 12:34:36 -04:00
parent 4bb20c9b71
commit cbaf4c4b64
12390 changed files with 318560 additions and 27248 deletions
Download patch file Download diff file Expand all files Collapse all files

391
Task/K-d-tree/COBOL/k-d-tree.cob Normal file
View file

@ -0,0 +1,391 @@
IDENTIFICATION DIVISION.
PROGRAM-ID. KD-TREE-TEST.
DATA DIVISION.
WORKING-STORAGE SECTION.
01 KD-TREE-STATE.
05 NUM-DIMENSIONS PIC 9(4) COMP.
05 TOTAL-NODES PIC 9(6) COMP.
05 TREE-ROOT PIC 9(6) COMP.
05 BEST-NODE PIC 9(6) COMP.
05 BEST-DISTANCE COMP-2.
05 VISITED-COUNT PIC 9(6) COMP.
05 MAX-NODES PIC 9(6) COMP VALUE 1005.
01 KD-NODES-DATA.
05 NODES OCCURS 1005 TIMES.
10 COORDS COMP-2 OCCURS 3 TIMES.
10 LEFT-CHILD PIC 9(6) COMP.
10 RIGHT-CHILD PIC 9(6) COMP.
01 KD-PTRS-DATA.
05 PTRS OCCURS 1005 TIMES.
10 NODE-PTR PIC 9(6) COMP.
01 TARGET-NODE-DATA.
05 TARGET-COORDS COMP-2 OCCURS 3 TIMES.
01 I PIC 9(6) COMP.
01 BEGIN-IDX PIC 9(6) COMP.
01 END-IDX PIC 9(6) COMP.
01 DIM-IDX PIC 9(4) COMP.
01 DISP-COORD-1 PIC -Z9.9(6).
01 DISP-COORD-2 PIC -Z9.9(6).
01 DISP-COORD-3 PIC -Z9.9(6).
01 TMP-DIST COMP-2.
01 DISP-DIST PIC -Z9.9(6).
01 DISP-VISITED PIC Z(5)9.
01 L-POINTS PIC 9(6) COMP.
01 SEED PIC 9(9) COMP.
01 DUMMY-FLOAT COMP-2.
PROCEDURE DIVISION.
MAIN-PROCEDURE.
MOVE FUNCTION CURRENT-DATE(11:6) TO SEED
COMPUTE DUMMY-FLOAT = FUNCTION RANDOM(SEED)
PERFORM TEST-WIKIPEDIA
DISPLAY " "
MOVE 100 TO L-POINTS
PERFORM TEST-RANDOM
DISPLAY " "
MOVE 1000 TO L-POINTS
PERFORM TEST-RANDOM
STOP RUN.
TEST-WIKIPEDIA.
INITIALIZE KD-NODES-DATA
MOVE 2 TO NUM-DIMENSIONS
MOVE 6 TO TOTAL-NODES
MOVE 0 TO TREE-ROOT
MOVE 0 TO BEST-NODE
MOVE 0 TO VISITED-COUNT
MOVE 2 TO COORDS(1, 1) MOVE 3 TO COORDS(1, 2)
MOVE 5 TO COORDS(2, 1) MOVE 4 TO COORDS(2, 2)
MOVE 9 TO COORDS(3, 1) MOVE 6 TO COORDS(3, 2)
MOVE 4 TO COORDS(4, 1) MOVE 7 TO COORDS(4, 2)
MOVE 8 TO COORDS(5, 1) MOVE 1 TO COORDS(5, 2)
MOVE 7 TO COORDS(6, 1) MOVE 2 TO COORDS(6, 2)
PERFORM VARYING I FROM 1 BY 1 UNTIL I > 6
MOVE I TO NODE-PTR OF PTRS(I)
END-PERFORM
MOVE 1 TO BEGIN-IDX
MOVE 7 TO END-IDX
MOVE 1 TO DIM-IDX
CALL "MAKE-TREE" USING KD-TREE-STATE KD-NODES-DATA KD-PTRS-DATA
BEGIN-IDX END-IDX DIM-IDX TREE-ROOT
MOVE 9 TO TARGET-COORDS(1)
MOVE 2 TO TARGET-COORDS(2)
MOVE 0 TO TARGET-COORDS(3)
MOVE 1 TO DIM-IDX
CALL "FIND-NEAREST" USING KD-TREE-STATE KD-NODES-DATA TARGET-NODE-DATA
TREE-ROOT DIM-IDX
*> DISPLAY "Tree dump:"
*> PERFORM VARYING I FROM 1 BY 1 UNTIL I > 6
*> DISPLAY "Node " I ": (" FUNCTION TRIM(COORDS(I, 1)) ", " FUNCTION TRIM(COORDS(I, 2)) ") "
*> "L=" LEFT-CHILD(I) " R=" RIGHT-CHILD(I)
*> END-PERFORM
DISPLAY "Wikipedia example data:"
COMPUTE TMP-DIST = FUNCTION SQRT(BEST-DISTANCE)
MOVE COORDS(BEST-NODE, 1) TO DISP-COORD-1
MOVE COORDS(BEST-NODE, 2) TO DISP-COORD-2
DISPLAY "nearest point: (" FUNCTION TRIM(DISP-COORD-1) ", "
FUNCTION TRIM(DISP-COORD-2) ")"
MOVE TMP-DIST TO DISP-DIST
DISPLAY "distance: " FUNCTION TRIM(DISP-DIST)
MOVE VISITED-COUNT TO DISP-VISITED
DISPLAY "nodes visited: " FUNCTION TRIM(DISP-VISITED).
TEST-RANDOM.
INITIALIZE KD-NODES-DATA
MOVE 3 TO NUM-DIMENSIONS
MOVE L-POINTS TO TOTAL-NODES
MOVE 0 TO TREE-ROOT
MOVE 0 TO BEST-NODE
MOVE 0 TO VISITED-COUNT
PERFORM VARYING I FROM 1 BY 1 UNTIL I > L-POINTS
COMPUTE COORDS(I, 1) = FUNCTION RANDOM
COMPUTE COORDS(I, 2) = FUNCTION RANDOM
COMPUTE COORDS(I, 3) = FUNCTION RANDOM
MOVE I TO NODE-PTR OF PTRS(I)
END-PERFORM
MOVE 1 TO BEGIN-IDX
COMPUTE END-IDX = L-POINTS + 1
MOVE 1 TO DIM-IDX
CALL "MAKE-TREE" USING KD-TREE-STATE KD-NODES-DATA KD-PTRS-DATA
BEGIN-IDX END-IDX DIM-IDX TREE-ROOT
COMPUTE TARGET-COORDS(1) = FUNCTION RANDOM
COMPUTE TARGET-COORDS(2) = FUNCTION RANDOM
COMPUTE TARGET-COORDS(3) = FUNCTION RANDOM
MOVE 1 TO DIM-IDX
CALL "FIND-NEAREST" USING KD-TREE-STATE KD-NODES-DATA TARGET-NODE-DATA
TREE-ROOT DIM-IDX
DISPLAY "Random data (" L-POINTS " points):"
MOVE TARGET-COORDS(1) TO DISP-COORD-1
MOVE TARGET-COORDS(2) TO DISP-COORD-2
MOVE TARGET-COORDS(3) TO DISP-COORD-3
DISPLAY "target: (" FUNCTION TRIM(DISP-COORD-1) ", "
FUNCTION TRIM(DISP-COORD-2) ", "
FUNCTION TRIM(DISP-COORD-3) ")"
COMPUTE TMP-DIST = FUNCTION SQRT(BEST-DISTANCE)
MOVE COORDS(BEST-NODE, 1) TO DISP-COORD-1
MOVE COORDS(BEST-NODE, 2) TO DISP-COORD-2
MOVE COORDS(BEST-NODE, 3) TO DISP-COORD-3
DISPLAY "nearest point: (" FUNCTION TRIM(DISP-COORD-1) ", "
FUNCTION TRIM(DISP-COORD-2) ", "
FUNCTION TRIM(DISP-COORD-3) ")"
MOVE TMP-DIST TO DISP-DIST
DISPLAY "distance: " FUNCTION TRIM(DISP-DIST)
MOVE VISITED-COUNT TO DISP-VISITED
DISPLAY "nodes visited: " FUNCTION TRIM(DISP-VISITED).
END PROGRAM KD-TREE-TEST.
IDENTIFICATION DIVISION.
PROGRAM-ID. MAKE-TREE IS RECURSIVE.
DATA DIVISION.
LOCAL-STORAGE SECTION.
01 QS-RIGHT PIC 9(6) COMP.
01 N-IDX PIC 9(6) COMP.
01 NEXT-DIM PIC 9(4) COMP.
01 N-PLUS-1 PIC 9(6) COMP.
LINKAGE SECTION.
01 KD-TREE-STATE.
05 NUM-DIMENSIONS PIC 9(4) COMP.
05 TOTAL-NODES PIC 9(6) COMP.
05 TREE-ROOT PIC 9(6) COMP.
05 BEST-NODE PIC 9(6) COMP.
05 BEST-DISTANCE COMP-2.
05 VISITED-COUNT PIC 9(6) COMP.
05 MAX-NODES PIC 9(6) COMP.
01 KD-NODES-DATA.
05 NODES OCCURS 1005 TIMES.
10 COORDS COMP-2 OCCURS 3 TIMES.
10 LEFT-CHILD PIC 9(6) COMP.
10 RIGHT-CHILD PIC 9(6) COMP.
01 KD-PTRS-DATA.
05 PTRS OCCURS 1005 TIMES.
10 NODE-PTR PIC 9(6) COMP.
01 BEGIN-IDX PIC 9(6) COMP.
01 END-IDX PIC 9(6) COMP.
01 DIM-IDX PIC 9(4) COMP.
01 RET-NODE PIC 9(6) COMP.
PROCEDURE DIVISION USING KD-TREE-STATE KD-NODES-DATA KD-PTRS-DATA
BEGIN-IDX END-IDX DIM-IDX RET-NODE.
IF END-IDX <= BEGIN-IDX
MOVE 0 TO RET-NODE
GOBACK
END-IF
COMPUTE N-IDX = BEGIN-IDX + (END-IDX - BEGIN-IDX) / 2
COMPUTE QS-RIGHT = END-IDX - 1
CALL "QUICK-SELECT" USING KD-TREE-STATE KD-NODES-DATA KD-PTRS-DATA
BEGIN-IDX QS-RIGHT N-IDX DIM-IDX RET-NODE
COMPUTE NEXT-DIM = FUNCTION MOD(DIM-IDX, NUM-DIMENSIONS) + 1
CALL "MAKE-TREE" USING KD-TREE-STATE KD-NODES-DATA KD-PTRS-DATA
BEGIN-IDX N-IDX NEXT-DIM LEFT-CHILD(RET-NODE)
COMPUTE N-PLUS-1 = N-IDX + 1
CALL "MAKE-TREE" USING KD-TREE-STATE KD-NODES-DATA KD-PTRS-DATA
N-PLUS-1 END-IDX NEXT-DIM RIGHT-CHILD(RET-NODE)
GOBACK.
END PROGRAM MAKE-TREE.
IDENTIFICATION DIVISION.
PROGRAM-ID. QUICK-SELECT.
DATA DIVISION.
WORKING-STORAGE SECTION.
01 RANGE-SIZE PIC 9(6) COMP.
01 PIVOT PIC 9(6) COMP.
01 PIVOT-VAL PIC 9(6) COMP.
01 TEMP PIC 9(6) COMP.
01 STORE-IDX PIC 9(6) COMP.
01 I PIC 9(6) COMP.
01 I-VAL PIC 9(6) COMP.
01 RET-PIVOT PIC 9(6) COMP.
LINKAGE SECTION.
01 KD-TREE-STATE.
05 NUM-DIMENSIONS PIC 9(4) COMP.
05 TOTAL-NODES PIC 9(6) COMP.
05 TREE-ROOT PIC 9(6) COMP.
05 BEST-NODE PIC 9(6) COMP.
05 BEST-DISTANCE COMP-2.
05 VISITED-COUNT PIC 9(6) COMP.
05 MAX-NODES PIC 9(6) COMP.
01 KD-NODES-DATA.
05 NODES OCCURS 1005 TIMES.
10 COORDS COMP-2 OCCURS 3 TIMES.
10 LEFT-CHILD PIC 9(6) COMP.
10 RIGHT-CHILD PIC 9(6) COMP.
01 KD-PTRS-DATA.
05 PTRS OCCURS 1005 TIMES.
10 NODE-PTR PIC 9(6) COMP.
01 L-LEFT PIC 9(6) COMP.
01 L-RIGHT PIC 9(6) COMP.
01 L-N PIC 9(6) COMP.
01 L-DIM PIC 9(4) COMP.
01 L-RET PIC 9(6) COMP.
PROCEDURE DIVISION USING KD-TREE-STATE KD-NODES-DATA KD-PTRS-DATA
L-LEFT L-RIGHT L-N L-DIM L-RET.
MAIN-LOOP.
PERFORM UNTIL 1 = 0
IF L-LEFT = L-RIGHT
MOVE NODE-PTR OF PTRS(L-LEFT) TO L-RET
GOBACK
END-IF
COMPUTE RANGE-SIZE = L-RIGHT - L-LEFT + 1
COMPUTE PIVOT = L-LEFT + FUNCTION INTEGER(RANGE-SIZE * FUNCTION RANDOM)
IF PIVOT > L-RIGHT
MOVE L-RIGHT TO PIVOT
END-IF
MOVE NODE-PTR OF PTRS(PIVOT) TO PIVOT-VAL
MOVE NODE-PTR OF PTRS(L-RIGHT) TO TEMP
MOVE PIVOT-VAL TO NODE-PTR OF PTRS(L-RIGHT)
MOVE TEMP TO NODE-PTR OF PTRS(PIVOT)
MOVE L-LEFT TO STORE-IDX
PERFORM VARYING I FROM L-LEFT BY 1 UNTIL I >= L-RIGHT
MOVE NODE-PTR OF PTRS(I) TO I-VAL
IF COORDS(I-VAL, L-DIM) < COORDS(PIVOT-VAL, L-DIM)
MOVE NODE-PTR OF PTRS(STORE-IDX) TO TEMP
MOVE I-VAL TO NODE-PTR OF PTRS(STORE-IDX)
MOVE TEMP TO NODE-PTR OF PTRS(I)
ADD 1 TO STORE-IDX
END-IF
END-PERFORM
MOVE NODE-PTR OF PTRS(STORE-IDX) TO TEMP
MOVE NODE-PTR OF PTRS(L-RIGHT) TO NODE-PTR OF PTRS(STORE-IDX)
MOVE TEMP TO NODE-PTR OF PTRS(L-RIGHT)
MOVE STORE-IDX TO RET-PIVOT
IF L-N = RET-PIVOT
MOVE NODE-PTR OF PTRS(L-N) TO L-RET
GOBACK
ELSE
IF L-N < RET-PIVOT
COMPUTE L-RIGHT = RET-PIVOT - 1
ELSE
COMPUTE L-LEFT = RET-PIVOT + 1
END-IF
END-IF
END-PERFORM.
END PROGRAM QUICK-SELECT.
IDENTIFICATION DIVISION.
PROGRAM-ID. FIND-NEAREST IS RECURSIVE.
DATA DIVISION.
LOCAL-STORAGE SECTION.
01 I PIC 9(4) COMP.
01 DIFF COMP-2.
01 DIST-SQ COMP-2.
01 DX COMP-2.
01 DX-SQ COMP-2.
01 NEXT-DIM PIC 9(4) COMP.
01 FIRST-CHILD PIC 9(6) COMP.
01 SECOND-CHILD PIC 9(6) COMP.
LINKAGE SECTION.
01 KD-TREE-STATE.
05 NUM-DIMENSIONS PIC 9(4) COMP.
05 TOTAL-NODES PIC 9(6) COMP.
05 TREE-ROOT PIC 9(6) COMP.
05 BEST-NODE PIC 9(6) COMP.
05 BEST-DISTANCE COMP-2.
05 VISITED-COUNT PIC 9(6) COMP.
05 MAX-NODES PIC 9(6) COMP.
01 KD-NODES-DATA.
05 NODES OCCURS 1005 TIMES.
10 COORDS COMP-2 OCCURS 3 TIMES.
10 LEFT-CHILD PIC 9(6) COMP.
10 RIGHT-CHILD PIC 9(6) COMP.
01 TARGET-NODE-DATA.
05 TARGET-COORDS COMP-2 OCCURS 3 TIMES.
01 L-CURR PIC 9(6) COMP.
01 L-DIM PIC 9(4) COMP.
PROCEDURE DIVISION USING KD-TREE-STATE KD-NODES-DATA TARGET-NODE-DATA
L-CURR L-DIM.
IF L-CURR = 0
GOBACK
END-IF
ADD 1 TO VISITED-COUNT
COMPUTE DIST-SQ = 0
PERFORM VARYING I FROM 1 BY 1 UNTIL I > NUM-DIMENSIONS
COMPUTE DIFF = TARGET-COORDS(I) - COORDS(L-CURR, I)
COMPUTE DIST-SQ = DIST-SQ + DIFF * DIFF
END-PERFORM
IF BEST-NODE = 0 OR DIST-SQ < BEST-DISTANCE
MOVE DIST-SQ TO BEST-DISTANCE
MOVE L-CURR TO BEST-NODE
END-IF
IF BEST-DISTANCE = 0
GOBACK
END-IF
COMPUTE DX = COORDS(L-CURR, L-DIM) - TARGET-COORDS(L-DIM)
COMPUTE NEXT-DIM = L-DIM + 1
IF NEXT-DIM > NUM-DIMENSIONS
MOVE 1 TO NEXT-DIM
END-IF
IF DX > 0
MOVE LEFT-CHILD(L-CURR) TO FIRST-CHILD
MOVE RIGHT-CHILD(L-CURR) TO SECOND-CHILD
ELSE
MOVE RIGHT-CHILD(L-CURR) TO FIRST-CHILD
MOVE LEFT-CHILD(L-CURR) TO SECOND-CHILD
END-IF
CALL "FIND-NEAREST" USING KD-TREE-STATE KD-NODES-DATA TARGET-NODE-DATA
FIRST-CHILD NEXT-DIM
COMPUTE DX-SQ = DX * DX
IF DX-SQ < BEST-DISTANCE
CALL "FIND-NEAREST" USING KD-TREE-STATE KD-NODES-DATA TARGET-NODE-DATA
SECOND-CHILD NEXT-DIM
END-IF
GOBACK.
END PROGRAM FIND-NEAREST.

552
Task/K-d-tree/Fortran/k-d-tree.f Normal file
View file

@ -0,0 +1,552 @@
! ==============================================================================
! 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