bktree/BKTree.hs at master · i-made/bktree · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
{--

An implementation of Burkhard-Keller trees.  These allow fast lookup
of words within a certain distance of a query word.

The implementation is based on the description in this article:

http://blog.notdot.net/2007/4/Damn-Cool-Algorithms-Part-1-BK-Trees

- Adam Hupp <adam@hupp.org>

--}

module BKTree ( queryTree,
                mkTree,
                levenshtein,
                main
              ) where


import List
import qualified Data.Map as Map

type Children a = Map.Map Int (TreeNode a)

data TreeNode a = TreeNode a (Children a) deriving Show

data Tree a = Tree (DistFunc a) (TreeNode a)

type DistFunc a = a -> a -> Int


addWord :: DistFunc a -> TreeNode a -> a -> TreeNode a
addWord distfn (TreeNode word kids) w =
    let distance = distfn word w
        newSubTree = if Map.member distance kids then
                         addWord distfn (kids Map.! distance) w
                     else
                         TreeNode w Map.empty
        childUpdate = Map.insert distance newSubTree kids
    in
      TreeNode word childUpdate


queryTreeNode :: DistFunc a -> a -> Int -> TreeNode a -> [(Int, a)]
queryTreeNode distfunc qword n (TreeNode word kids)  =
    let dist = distfunc word qword
        range = [dist-n..dist+n+1]
        kidsInRange = mapSelect kids range
        childResults = concatMap (queryTreeNode distfunc qword n) kidsInRange
    in if dist <= n then
           (dist, word) : childResults
       else
           childResults

compareBy :: Ord b => (a -> b) -> a -> a -> Ordering
compareBy by lhs rhs = compare (by lhs) (by rhs)

mapSelect :: Ord k => Map.Map k v -> [k] -> [v]
mapSelect mp (x:xs) =
    if Map.member x mp then
        (mp Map.! x ) : mapSelect mp xs
    else mapSelect mp xs
mapSelect _ [] = []

mkTree :: DistFunc a -> [a] -> Tree a
mkTree distfn (w:ws) =
    let root = TreeNode w Map.empty
        tree = foldl (addWord distfn) root ws
    in
      Tree distfn tree


queryTree :: Tree a -> a -> Int -> [(Int,a)]
queryTree (Tree distfn root) qword n =
    sortBy (compareBy fst) $ queryTreeNode distfn qword n root

main = do dictFile <- readFile "/usr/share/dict/american-english"
          words <- return $ lines dictFile
          tree <- return $ mkTree levenshtein words
          putStrLn (show (queryTree tree "the" 2))

-- http://www.cse.unsw.edu.au/~dons/code/lambdabot/Lib/Util.hs
-- | Levenshtein edit-distance algorithm
-- Translated from an Erlang version by Fredrik Svensson and Adam Lindberg
--
levenshtein :: String -> String -> Int
levenshtein [] [] = 0
levenshtein s  [] = length s
levenshtein [] s  = length s
levenshtein s  t  = lvn s t [0..length t] 1

lvn :: String -> String -> [Int] -> Int -> Int
lvn [] _ dl _ = last dl
lvn (s:ss) t dl n = lvn ss t (lvn' t dl s [n] n) (n + 1)

lvn' :: String -> [Int] -> Char -> [Int] -> Int -> [Int]
lvn' [] _ _ ndl _ = ndl
lvn' (t:ts) (dlh:dlt) c ndl ld | length dlt > 0 = lvn' ts dlt c (ndl ++ [m]) m
    where
        m = foldl1 min [ld + 1, head dlt + 1, dlh + (dif t c)]
lvn' _ _ _ _  _  = error "levenshtein, ran out of numbers"

dif :: Char -> Char -> Int
dif = (fromEnum .) . (/=)