Parasitemia/Parasitemia/KdTree.fs

   1 module KdTree
   2
   3 open System
   4
   5 type I2DCoords =
   6     abstract X : float32
   7     abstract Y : float32
   8
   9 // Compare 'e1' and 'e2' by X.
  10 let cmpX (e1: I2DCoords) (e2: I2DCoords) : int =
  11     match e1.X.CompareTo(e2.X) with
  12     | 0 -> match e1.Y.CompareTo(e2.Y) with
  13            | 0 -> e1.GetHashCode().CompareTo(e2.GetHashCode())
  14            | v -> v
  15     | v -> v
  16
  17 // Compare 'e1' and 'e2' by Y.
  18 let cmpY (e1: I2DCoords) (e2: I2DCoords) : int =
  19     match e1.Y.CompareTo(e2.Y) with
  20     | 0 -> match e1.X.CompareTo(e2.X) with
  21            | 0 -> e1.GetHashCode().CompareTo(e2.GetHashCode())
  22            | v -> v
  23     | v -> v
  24
  25 type Region = { minX: float32; maxX: float32; minY: float32; maxY: float32 } with
  26     member this.Contains px py : bool =
  27         px >= this.minX && px <= this.maxX &&
  28         py >= this.minY && py <= this.maxY
  29
  30     member this.IsSub otherRegion : bool =
  31         this.minX >= otherRegion.minX && this.maxX <= otherRegion.maxX &&
  32         this.minY >= otherRegion.minY && this.maxY <= otherRegion.maxY
  33
  34     member this.Intersects otherRegion : bool =
  35         this.minX < otherRegion.maxX && this.maxX >= otherRegion.minX &&
  36         this.minY < otherRegion.maxY && this.maxY >= otherRegion.minY
  37
  38 type Tree<'a when 'a :> I2DCoords> =
  39     | Node of float32 * Tree<'a> * Tree<'a>
  40     | Leaf of 'a
  41
  42     static member BuildTree (l: 'a list) : Tree<'a> =
  43         let xSorted = List.toArray l
  44         let ySorted = List.toArray l
  45         Array.sortInPlaceWith cmpX xSorted
  46         Array.sortInPlaceWith cmpY ySorted
  47
  48         let rec buildTreeFromSortedArray (pXSorted: 'a[]) (pYSorted: 'a[]) (depth: int) : Tree<'a> =
  49             if pXSorted.Length = 1
  50             then
  51                 Leaf pXSorted.[0]
  52             else
  53                 if depth % 2 = 1 // 'depth' is odd -> vertical splitting else horizontal splitting.
  54                 then
  55                     let leftX, rightX = Array.splitAt ((pXSorted.Length + 1) / 2) pXSorted
  56                     let splitElement = Array.last leftX
  57                     let leftY, rightY = Array.partition (fun (e: 'a) -> cmpX e splitElement <= 0) pYSorted // FIXME: Maybe this operation can be optimized.
  58                     Node (splitElement.X, buildTreeFromSortedArray leftX leftY (depth + 1), buildTreeFromSortedArray rightX rightY (depth + 1))
  59                 else
  60                     let downY, upY = Array.splitAt ((pYSorted.Length + 1) / 2) pYSorted
  61                     let splitElement = Array.last downY
  62                     let downX, upX = Array.partition (fun (e: 'a) -> cmpY e splitElement <= 0) pXSorted // FIXME: Maybe this operation can be optimized.
  63                     Node (splitElement.Y, buildTreeFromSortedArray downX downY (depth + 1), buildTreeFromSortedArray upX upY (depth + 1))
  64
  65         buildTreeFromSortedArray xSorted ySorted 1
  66
  67     member this.Search (searchRegion: Region) : 'a list =
  68         let rec valuesFrom (tree: Tree<'a>) : 'a list =
  69             match tree with
  70             | Leaf v -> [v]
  71             | Node (_, part1, part2) -> (valuesFrom part1) @ (valuesFrom part2)
  72
  73         let rec searchWithRegion (tree: Tree<'a>) (currentRegion: Region) (depth: int) : 'a list =
  74             match tree with
  75             | Leaf v -> if searchRegion.Contains v.X v.Y then [v] else []
  76             | Node (splitValue, part1, part2) ->
  77                 let valuesInRegion (region: Region) (treeRegion: Tree<'a>) =
  78                     if region.IsSub searchRegion
  79                     then
  80                         valuesFrom treeRegion
  81                     elif region.Intersects searchRegion
  82                     then
  83                         searchWithRegion treeRegion region (depth + 1)
  84                     else
  85                         []
  86
  87                 if depth % 2 = 1 // Vertical splitting.
  88                 then
  89                     let leftRegion = { currentRegion with maxX = splitValue }
  90                     let rightRegion = { currentRegion with minX = splitValue }
  91                     (valuesInRegion leftRegion part1) @ (valuesInRegion rightRegion part2)
  92                 else // Horizontal splitting.
  93                     let downRegion = { currentRegion with maxY = splitValue }
  94                     let upRegion = { currentRegion with minY = splitValue }
  95                     (valuesInRegion downRegion part1) @ (valuesInRegion upRegion part2)
  96
  97         searchWithRegion this { minX = Single.MinValue; maxX = Single.MaxValue; minY = Single.MinValue; maxY = Single.MaxValue } 1
  98
  99 ///// Tests. TODO: to put in a unit test.
 100
 101 type Point (x: float32, y: float32) =
 102     interface I2DCoords with
 103         member this.X = x
 104         member this.Y = y
 105
 106     override this.ToString () =
 107         sprintf "(%.1f, %.1f)" x y
 108
 109 // TODO: test with identical X or Y coords
 110 let test () =
 111     let pts = [
 112         Point(1.0f, 1.0f)
 113         Point(2.0f, 2.0f)
 114         Point(1.5f, 3.6f)
 115         Point(3.0f, 3.2f)
 116         Point(4.0f, 4.0f)
 117         Point(3.5f, 1.5f)
 118         Point(2.5f, 0.5f) ]
 119
 120     let tree = Tree.BuildTree pts
 121     Utils.dprintfn "Tree: %A" tree
 122
 123     let s1 = tree.Search { minX = 0.0f; maxX = 5.0f; minY = 0.0f; maxY = 5.0f } // All points.
 124     Utils.dprintfn "s1: %A" s1
 125
 126     let s2 = tree.Search { minX = 2.8f; maxX = 4.5f; minY = 3.0f; maxY = 4.5f }
 127     Utils.dprintfn "s2: %A" s2
 128
 129     let s3 = tree.Search { minX = 2.0f; maxX = 2.0f; minY = 2.0f; maxY = 2.0f }
 130     Utils.dprintfn "s3: %A" s3
 131
 132 let test2 () =
 133     let pts = [
 134         Point(1.0f, 1.0f)
 135         Point(1.0f, 2.0f)
 136         Point(1.0f, 3.0f) ]
 137
 138     let tree = Tree.BuildTree pts
 139     Utils.dprintfn "Tree: %A" tree
 140
 141     let s1 = tree.Search { minX = 1.0f; maxX = 1.0f; minY = 1.0f; maxY = 1.0f }
 142     Utils.dprintfn "s1: %A" s1
 143
 144     let s2 = tree.Search { minX = 1.0f; maxX = 1.0f; minY = 2.0f; maxY = 2.0f }
 145     Utils.dprintfn "s2: %A" s2
 146
 147     // This case result is wrong: FIXME
 148     let s3 = tree.Search { minX = 1.0f; maxX = 1.0f; minY = 3.0f; maxY = 3.0f }
 149     Utils.dprintfn "s3: %A" s3
 150
 151     let s4 = tree.Search { minX = 0.0f; maxX = 2.0f; minY = 0.0f; maxY = 4.0f }
 152     Utils.dprintfn "s4: %A" s4
 153