diff --git a/Data/BKTree.hs b/Data/BKTree.hs new file mode 100644 --- /dev/null +++ b/Data/BKTree.hs @@ -0,0 +1,69 @@ +-- | Implementation of a BK-tree: +module Data.BKTree + ( -- * Types + Distance + , BKTree + -- * Operations + , empty + , insert + , query + ) where + +import Data.IntMap (IntMap) +import qualified Data.IntMap as IntMap + +type Distance s = s -> s -> Int + +data BKTree s = + BKTree !(BK s) (Distance s) + +data BK s + = EmptyBK + | BK !s !(IntMap (BK s)) + +narrow :: Int -> Int -> IntMap a -> IntMap a +narrow n m im | n == m = IntMap.fromList (maybe [] (\v -> [(n, v)]) (IntMap.lookup n im)) +narrow n m im | otherwise = insMaybe m res pr + where + (_, pl, res0) = IntMap.splitLookup n im + (res, pr, _) = IntMap.splitLookup m (insMaybe n res0 pl) + insMaybe k im' = maybe im' (\v -> IntMap.insert k v im') + +empty :: Distance s -- ^ The distance function \"d\" must be a metric on \"s\" + -- (): + -- + -- * d x y >= 0 + -- + -- * d x y == 0 iff x == y + -- + -- * d x y == d y x + -- + -- * d x z <= d x y + d y z + -> BKTree s +empty = BKTree EmptyBK + +insert :: s -> BKTree s -> BKTree s +insert s (BKTree bk f) = BKTree (insert' s f bk) f + +insert' :: s -> Distance s -> BK s -> BK s +insert' s _ EmptyBK = BK s IntMap.empty +insert' s f bk@(BK s' bks) + | dist == 0 = bk + | otherwise = BK s' $ flip (IntMap.insert dist) bks $ + maybe (insert' s f EmptyBK) (insert' s f) (IntMap.lookup dist bks) + where dist = f s s' + +query :: Int -- ^ The maximum distance to search for. + -> s -> BKTree s + -> [(s, Int)] -- ^ All the words with a distance less than the + -- one specified, and their respective + -- distances. +query maxd s (BKTree bk f) = query' maxd s f bk + +query' :: Int -> s -> Distance s -> BK s -> [(s, Int)] +query' _ _ _ EmptyBK = [] +query' maxd s f (BK s' bks) = match ++ concatMap (query' maxd s f) children + where + dist = f s s' + match = if (dist <= maxd) then [(s', dist)] else [] + children = IntMap.elems $ narrow (max (dist - maxd) 0) (dist + maxd) bks diff --git a/Language/Distance/Search/BK.hs b/Language/Distance/Search/BK.hs --- a/Language/Distance/Search/BK.hs +++ b/Language/Distance/Search/BK.hs @@ -6,8 +6,7 @@ -- . It performs reasonably, and it -- scales decently as the query distance increases. Moreover the data -- structure can work on any instance of 'EditDistance', or in fact any metric --- space (although no interface for that purpose is defined): --- . +-- space - a generic interface is provided in 'Data.BKTree'. -- -- However, for very short distances (less than 3), -- 'Language.Distance.Search.TST' is faster. @@ -22,47 +21,27 @@ , damerauLevenshtein ) where -import Data.IntMap (IntMap) -import qualified Data.IntMap as IntMap +import Control.Arrow (second) import Data.ListLike (ListLike) +import qualified Data.BKTree as BKTree import Language.Distance (EditDistance (..), Levenshtein, DamerauLevenshtein) import Language.Distance.Internal -data BKTree full algo - = EmptyBK - | BKTree !full !(IntMap (BKTree full algo)) - -narrow :: Int -> Int -> IntMap a -> IntMap a -narrow n m im | n == m = IntMap.fromList (maybe [] (\v -> [(n, v)]) (IntMap.lookup n im)) -narrow n m im | otherwise = insMaybe m res pr - where - (_, pl, res0) = IntMap.splitLookup n im - (res, pr, _) = IntMap.splitLookup m (insMaybe n res0 pl) - insMaybe k im' = maybe im' (\v -> IntMap.insert k v im') +newtype BKTree full algo = BKTree (BKTree.BKTree full) -empty :: BKTree full algo -empty = EmptyBK +empty :: forall full sym algo. (EditDistance sym algo, ListLike full sym) + => BKTree full algo +empty = BKTree (BKTree.empty (\s s' -> getDistance (distance s s' :: Distance algo))) -insert :: forall full sym algo. (Eq sym, EditDistance sym algo, ListLike full sym) +insert :: (EditDistance sym algo, ListLike full sym) => full -> BKTree full algo -> BKTree full algo -insert str EmptyBK = BKTree str IntMap.empty -insert str bk@(BKTree str' bks) - | dist == 0 = bk - | otherwise = BKTree str' $ flip (IntMap.insert dist) bks $ - maybe (insert str EmptyBK) (insert str) (IntMap.lookup dist bks) - where dist = getDistance (distance str str' :: Distance algo) +insert s (BKTree bk) = BKTree (BKTree.insert s bk) -query :: forall full sym algo. (ListLike full sym, EditDistance sym algo) +query :: (ListLike full sym, EditDistance sym algo) => Int -> full -> BKTree full algo -> [(full, Distance algo)] -query _ _ EmptyBK = [] -query maxd str (BKTree str' bks) = match ++ concatMap (query maxd str) children - where - dist = distance str str' :: Distance algo - intDist = getDistance dist - match = if (intDist <= maxd) then [(str', dist)] else [] - children = IntMap.elems $ narrow (max (intDist - maxd) 0) (intDist + maxd) bks +query maxd s (BKTree bk) = map (second Distance) $ BKTree.query maxd s bk levenshtein :: (ListLike full sym, EditDistance sym Levenshtein) => Int -> full -> BKTree full Levenshtein -> [(full, Distance Levenshtein)] diff --git a/language-spelling.cabal b/language-spelling.cabal --- a/language-spelling.cabal +++ b/language-spelling.cabal @@ -1,6 +1,6 @@ Cabal-Version: >= 1.8 Name: language-spelling -Version: 0.1.2 +Version: 0.2 Author: Francesco Mazzoli (f@mazzo.li) Maintainer: Francesco Mazzoli (f@mazzo.li) Build-Type: Simple @@ -55,15 +55,16 @@ Exposed-Modules: Language.Distance, Language.Distance.Internal, Language.Distance.Search, - Language.Distance.Search.Class, Language.Distance.Search.BK, Language.Distance.Search.TST, Language.Phonetic, Language.Phonetic.Internal, + Data.BKTree, Data.TST, Data.TSTSet - Other-Modules: Language.Phonetic.Soundex, + Other-Modules: Language.Distance.Search.Class, + Language.Phonetic.Soundex, Language.Phonetic.Encoder Test-Suite benchmarks