Implementacja kd-drzewa.

kdtree.hs

import Data.List

data Tree a = EmptyNode 
  | Node { value:: a, low :: Tree a, high :: Tree a } deriving (Show)

type Point = [Int]

cmpAlong :: Int -> Point -> Point -> Ordering
cmpAlong ax p1 p2
  | p1 !! ax > p2 !! ax = GT
  | p1 !! ax < p2 !! ax = LT
  | otherwise           = EQ

partitionAlong :: Int -> [Point] -> ([Point], Point, [Point])
partitionAlong ax points = (low, median, high) where
  sorted = sortBy (cmpAlong ax) points
  median = sorted !! (length points `div` 2)
  mAx = median !! ax
  (low, (_ : high)) = partition (\x -> x !! ax < mAx) sorted

inRange value range = x_min <= x && x <= x_max && y_min <= y && y <= y_max where
  [x, y] = value
  ([x_min, y_min], [x_max, y_max]) = range

buildKdTree :: Int -> [Point] -> Tree Point
buildKdTree k points = build 0 points where 
  build _ [ ]        = EmptyNode
  build depth points = Node {
    value = median,
    low = build (depth + 1) low,
    high = build (depth + 1) high
  } where
    ax = depth `mod` k
    (low, median, high) = partitionAlong ax points

query :: Int -> Tree Point -> (Point, Point) -> [Point]
query k root range = query 0 root where
  (min, max) = range
  query _ EmptyNode = []
  query depth Node { value = p, low = low, high = high }
    | p !! ax < min !! ax = query (depth + 1) high
    | p !! ax > max !! ax = query (depth + 1) low
    | p `inRange` range       = p : query (depth + 1) low ++ query (depth + 1) high
    | otherwise               = query (depth + 1) low ++ query (depth + 1) high
    where 
      ax = depth `mod` k