Skip to content

Instantly share code, notes, and snippets.

Created March 7, 2010 19:25
Show Gist options
  • Save anonymous/324580 to your computer and use it in GitHub Desktop.
Save anonymous/324580 to your computer and use it in GitHub Desktop.
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
cut (x,y) = fst . IM.split (y+1) . snd . IM.split (x-1)
lookup :: MetricSpace a => a -> BKTree a -> Maybe a
lookup x = listToMaybe . query x 0
fromList :: MetricSpace a => [a] -> BKTree a
fromList = foldr insert empty
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment