Parasitemia/ParasitemiaCore/KdTree.fs

   1 module ParasitemiaCore.KdTree
   2
   3 open System
   4 open System.Collections.Generic
   5
   6 type I2DCoords =
   7     abstract X : float32
   8     abstract Y : float32
   9
  10 // Compare 'e1' and 'e2' by X.
  11 let cmpX (e1 : I2DCoords) (e2 : I2DCoords) : int =
  12     match e1.X.CompareTo e2.X with
  13     | 0 -> match e1.Y.CompareTo e2.Y with
  14            | 0 -> e1.GetHashCode().CompareTo (e2.GetHashCode ())
  15            | v -> v
  16     | v -> v
  17
  18 // Compare 'e1' and 'e2' by Y.
  19 let cmpY (e1 : I2DCoords) (e2 : I2DCoords) : int =
  20     match e1.Y.CompareTo e2.Y with
  21     | 0 -> match e1.X.CompareTo e2.X with
  22            | 0 -> e1.GetHashCode().CompareTo (e2.GetHashCode ())
  23            | v -> v
  24     | v -> v
  25
  26 type Region = { minX : float32; maxX : float32; minY : float32; maxY : float32 } with
  27     member this.Contains px py : bool =
  28         px >= this.minX && px <= this.maxX &&
  29         py >= this.minY && py <= this.maxY
  30
  31     member this.IsSub otherRegion : bool =
  32         this.minX >= otherRegion.minX && this.maxX <= otherRegion.maxX &&
  33         this.minY >= otherRegion.minY && this.maxY <= otherRegion.maxY
  34
  35     member this.Intersects otherRegion : bool =
  36         this.minX < otherRegion.maxX && this.maxX >= otherRegion.minX &&
  37         this.minY < otherRegion.maxY && this.maxY >= otherRegion.minY
  38
  39 type Tree<'a when 'a :> I2DCoords> =
  40     | Node of float32 * Tree<'a> * Tree<'a>
  41     | Leaf of 'a
  42
  43     static member BuildTree (l : 'a list) : Tree<'a> =
  44         let xSorted = List.toArray l
  45         let ySorted = List.toArray l
  46         Array.sortInPlaceWith cmpX xSorted
  47         Array.sortInPlaceWith cmpY ySorted
  48
  49         let rec buildTreeFromSortedArray (pXSorted : 'a[]) (pYSorted : 'a[]) (depth : int) : Tree<'a> =
  50             if pXSorted.Length = 1 then
  51                 Leaf pXSorted.[0]
  52             else
  53                 if depth % 2 = 1 then // 'depth' is odd -> vertical splitting else horizontal splitting.
  54                     let leftX, rightX = Array.splitAt ((pXSorted.Length + 1) / 2) pXSorted
  55                     let splitElement = Array.last leftX
  56                     let leftY, rightY = Array.partition (fun (e : 'a) -> cmpX e splitElement <= 0) pYSorted // FIXME: Maybe this operation can be optimized.
  57                     Node (splitElement.X, buildTreeFromSortedArray leftX leftY (depth + 1), buildTreeFromSortedArray rightX rightY (depth + 1))
  58                 else
  59                     let downY, upY = Array.splitAt ((pYSorted.Length + 1) / 2) pYSorted
  60                     let splitElement = Array.last downY
  61                     let downX, upX = Array.partition (fun (e : 'a) -> cmpY e splitElement <= 0) pXSorted // FIXME: Maybe this operation can be optimized.
  62                     Node (splitElement.Y, buildTreeFromSortedArray downX downY (depth + 1), buildTreeFromSortedArray upX upY (depth + 1))
  63
  64         buildTreeFromSortedArray xSorted ySorted 1
  65
  66     member this.Search (searchRegion : Region) : List<'a> =
  67         let result = List<'a> ()
  68         let rec valuesFrom (tree : Tree<'a>) =
  69             match tree with
  70             | Node (_, left, right) ->
  71                 valuesFrom right
  72                 valuesFrom left
  73             | Leaf v ->
  74                 result.Add v
  75
  76         let rec searchWithRegion (tree : Tree<'a>) (currentRegion : Region) (depth : int) =
  77             match tree with
  78             | Leaf v ->
  79                 if searchRegion.Contains v.X v.Y then
  80                     result.Add v
  81             | Node (splitValue, part1, part2) ->
  82                 let inline valuesInRegion (region : Region) (treeRegion : Tree<'a>) =
  83                     if region.IsSub searchRegion then
  84                         valuesFrom treeRegion
  85                     elif region.Intersects searchRegion then
  86                         searchWithRegion treeRegion region (depth + 1)
  87
  88                 if depth % 2 = 1 then // Vertical splitting.
  89                     valuesInRegion { currentRegion with maxX = splitValue } part1 // Left region.
  90                     valuesInRegion { currentRegion with minX = splitValue } part2 // Right region.
  91                 else // Horizontal splitting.
  92                     valuesInRegion { currentRegion with maxY = splitValue } part1 // Down region.
  93                     valuesInRegion { currentRegion with minY = splitValue } part2 // Up region.
  94
  95         searchWithRegion this { minX = Single.MinValue; maxX = Single.MaxValue; minY = Single.MinValue; maxY = Single.MaxValue } 1
  96         result