BKTree.hs

module Data.BKTree (
        BKTree,
        MetricSpace(..),

        empty,
        null,
        size,
        singleton,
        insert,
        query,
        lookup,

        fromList
  ) where

import Data.Maybe (listToMaybe)

import Data.IntMap (IntMap)
import qualified Data.IntMap as IM

import Prelude hiding (lookup,null)

class MetricSpace a where
  distance :: a -> a -> Int

-- Dummy instance for testing.
instance MetricSpace Int where
  distance x y = abs (y - x)

data BKTree a = Node a !(IntMap (BKTree a))
              | Empty
  deriving Show

empty :: MetricSpace a => BKTree a
empty = Empty

null :: MetricSpace a => BKTree a -> Bool
null Empty = True
null _     = False

size :: MetricSpace a => BKTree a -> Int
size Empty           = 0
size (Node _ childs) = 1 + (sum . map size . IM.elems $ childs)

singleton :: MetricSpace a => a -> BKTree a
singleton x = Node x IM.empty

insert :: MetricSpace a => a -> BKTree a -> BKTree a
insert x Empty            = singleton x
insert x (Node x0 childs) = Node x0 (IM.alter add d childs)
  where add = Just . maybe (singleton x) (insert x)
        d   = distance x x0

query :: MetricSpace a => a ->  Int -> BKTree a -> [a]
query _ _ Empty            = []
query x n (Node x0 childs) = if d <= n then x0 : go else go
  where go  = concatMap (query x n)
                        (IM.elems . inRange (d-n,d+n) $ childs)
        d   = distance x x0

lookup :: MetricSpace a => a -> BKTree a -> Maybe a
lookup x = listToMaybe . query x 0

fromList :: MetricSpace a => [a] -> BKTree a
fromList = foldr insert empty