edit distance algorithm in Haskell - performance tuning

Question

I\'m trying to implement the levenshtein distance (or edit distance) in Haskell, but its performance decreases rapidly when the string lenght increases.

I\'m still q

Answer 1

You need to memoize editDistance'. There are many ways of doing this, e.g., a recursively defined array.

Answer 2

People are recommending you to use generic memoization libraries, but for the simple task of defining Levenshtein distance plain dynamic programming is more than enough. A very simple polymorphic list-based implementation:

distance s t = 
    d !!(length s)!!(length t) 
    where d = [ [ dist m n | n <- [0..length t] ] | m <- [0..length s] ]
          dist i 0 = i
          dist 0 j = j
          dist i j = minimum [ d!!(i-1)!!j+1
                             , d!!i!!(j-1)+1
                             , d!!(i-1)!!(j-1) + (if s!!(i-1)==t!!(j-1) 
                                                  then 0 else 1) 
                             ]

Or if you need real speed on long sequences, you can use a mutable array:

import Data.Array
import qualified Data.Array.Unboxed as UA
import Data.Array.ST
import Control.Monad.ST


-- Mutable unboxed and immutable boxed arrays
distance :: Eq a => [a] -> [a] -> Int
distance s t = d UA.! (ls , lt)
    where s' = array (0,ls) [ (i,x) | (i,x) <- zip [0..] s ]
          t' = array (0,lt) [ (i,x) | (i,x) <- zip [0..] t ]
          ls = length s
          lt = length t
          (l,h) = ((0,0),(length s,length t))
          d = runSTUArray $ do
                m <- newArray (l,h) 0 
                for_ [0..ls] $ \i -> writeArray m (i,0) i
                for_ [0..lt] $ \j -> writeArray m (0,j) j
                for_ [1..lt] $ \j -> do
                              for_ [1..ls] $ \i -> do
                                  let c = if s'!(i-1)==t'! (j-1) 
                                          then 0 else 1
                                  x <- readArray m (i-1,j)
                                  y <- readArray m (i,j-1)
                                  z <- readArray m (i-1,j-1)
                                  writeArray m (i,j) $ minimum [x+1, y+1, z+c ]
                return m

for_ xs f =  mapM_ f xs

Answer 3

As already mentioned, memoization is what you need. In addition you are looking at edit distance from right to left, wich isn't very efficient with strings, and edit distance is the same regardless of direction. That is: editDistance (reverse a) (reverse b) == editDistance a b

For solving the memoization part there are very many libraries that can help you. In my example below I chose MemoTrie since it is quite easy to use and performs well here.

import Data.MemoTrie(memo2)

editDistance' del sub ins = memf
  where
   memf = memo2 f
   f s1     []     = ins * length s1
   f []     s2     = ins * length s2
   f (x:xs) (y:ys)
     | x == y  = memf xs ys
     | otherwise = minimum [ del + memf xs (y:ys),
                             sub + memf (x:xs) ys,
                             ins + memf xs ys]

As you can see all you need is to add the memoization. The rest is the same except that we start from the beginning of the list in staid of the end.

Answer 4

Ignoring that this is a bad algorithm (should be memoizing, I get to that second)...

Use O(1) Primitives and not O(n)

One problem is you use a whole bunch calls that are O(n) for lists (haskell lists are singly linked lists). A better data structure would give you O(1) operations, I used Vector:

import qualified Data.Vector as V

-- standard levenshtein distance between two lists
editDistance      :: Eq a => [a] -> [a] -> Int
editDistance s1 s2 = editDistance' 1 1 1 (V.fromList s1) (V.fromList s2)

-- weighted levenshtein distance
-- ins, sub and del are the costs for the various operations
editDistance'      :: Eq a => Int -> Int -> Int -> V.Vector a -> V.Vector a -> Int
editDistance' del sub ins s1 s2
  | V.null s2 = ins * V.length s1
  | V.null s1 = ins * V.length s2
  | V.last s1 == V.last s2 = editDistance' del sub ins (V.init s1) (V.init s2)
  | otherwise            = minimum [ editDistance' del sub ins s1 (V.init s2)        + del -- deletion 
                                   , editDistance' del sub ins (V.init s1) (V.init s2) + sub -- substitution
                                   , editDistance' del sub ins (V.init s1) s2        + ins -- insertion
                                   ]

The operations that are O(n) for lists include init, length, and last (though init is able to be lazy at least). All these operations are O(1) using Vector.

While real benchmarking should use Criterion, a quick and dirty benchmark:

str2 = replicate 15 'a' ++ replicate 25 'b'
str1 = replicate 20 'a' ++ replicate 20 'b'
main = print $ editDistance str1 str2

shows the vector version takes 0.09 seconds while strings take 1.6 seconds, so we saved about an order of magnitude without even looking at your editDistance algorithm.

Now what about memoizing results?

The bigger issue is obviously the need for memoization. I took this as an opportunity to learn the monad-memo package - my god is that awesome! For one extra constraint (you need Ord a), you get a memoization basically for no effort. The code:

import qualified Data.Vector as V
import Control.Monad.Memo

-- standard levenshtein distance between two lists
editDistance      :: (Eq a, Ord a) => [a] -> [a] -> Int
editDistance s1 s2 = startEvalMemo $ editDistance' (1, 1, 1, (V.fromList s1), (V.fromList s2))

-- weighted levenshtein distance
-- ins, sub and del are the costs for the various operations
editDistance' :: (MonadMemo (Int, Int, Int, V.Vector a, V.Vector a) Int m, Eq a) => (Int, Int, Int, V.Vector a, V.Vector a) -> m Int
editDistance' (del, sub, ins, s1, s2)
  | V.null s2 = return $ ins * V.length s1
  | V.null s1 = return $ ins * V.length s2
  | V.last s1 == V.last s2 = memo editDistance' (del, sub, ins, (V.init s1), (V.init s2))
  | otherwise = do
        r1 <- memo editDistance' (del, sub, ins, s1, (V.init s2))
        r2 <- memo editDistance' (del, sub, ins, (V.init s1), (V.init s2))
        r3 <- memo editDistance' (del, sub, ins, (V.init s1), s2)
        return $ minimum [ r1 + del -- deletion 
                         , r2 + sub -- substitution
                         , r3 + ins -- insertion
                                   ]

You see how the memoization needs a single "key" (see the MonadMemo class)? I packaged all the arguments as a big ugly tuple. It also needs one "value", which is your resulting Int. Then it's just plug and play using the "memo" function for the values you want to memoize.

For benchmarking I used a shorter, but larger-distance, string:

$ time ./so  # the memoized vector version
12

real    0m0.003s

$ time ./so3  # the non-memoized vector version
12

real    1m33.122s

Don't even think about running the non-memoized string version, I figure it would take around 15 minutes at a minimum. As for me, I now love monad-memo - thanks for the package Eduard!

EDIT: The difference between String and Vector isn't as much in the memoized version, but still grows to a factor of 2 when the distance gets to around 200, so still worth while.

EDIT: Perhaps I should explain why the bigger issue is "obviously" memoizing results. Well, if you look at the heart of the original algorithm:

 [ editDistance' ... s1          (V.init s2)  + del 
 , editDistance' ... (V.init s1) (V.init s2) + sub
 , editDistance' ... (V.init s1) s2          + ins]

It's quite clear a call of editDistance' s1 s2 results in 3 calls to editDistance'... each of which call editDistance' three more times... and three more time... and AHHH! Exponential explosion! Luckly most the calls are identical! for example (using --> for "calls" and eD for editDistance'):

eD s1 s2  --> eD s1 (init s2)             -- The parent
            , eD (init s1) s2
            , eD (init s1) (init s2)
eD (init s1) s2 --> eD (init s1) (init s2)         -- The first "child"
                  , eD (init (init s1)) s2
                  , eD (init (init s1)) (init s2) 
eD s1 (init s2) --> eD s1 (init (init s2))
                  , eD (init s1) (init s2)
                  , eD (init s1) (init (init s2))

Just by considering the parent and two immediate children we can see the call ed (init s1) (init s2) is done three times. The other child share calls with the parent too and all children share many calls with each other (and their children, cue Monty Python skit).

It would be a fun, perhaps instructive, exercise to make a runMemo like function that returns the number of cached results used.

Answer 5

I know there's already an editDistance implementation on Hackage, but I need it to work on lists of arbitrary tokens, not necessarily strings

Are there a finite number of tokens? I'd suggest you try simply devising a mapping from token to character. There are 10,646 characters at your disposal, after all.

Answer 6

This version is much quicker than those memorised versions, but still I would love to have it even quicker. Works fine with 100's character long strings. I was written with other distances in mind(change the init function and cost) , and use classical dynamic programming array trick. The long line could be converted into a separate function with top 'do', but I like this way.

import Data.Array.IO
import System.IO.Unsafe

editDistance = dist ini med

dist :: (Int -> Int -> Int) -> (a -> a -> Int ) -> [a] -> [a] -> Int
dist i f a b  = unsafePerformIO $ distM i f a b

-- easy to create other distances 
ini i 0 = i
ini 0 j = j
ini _ _ = 0
med a b = if a == b then 0 else 2


distM :: (Int -> Int -> Int) -> (a -> a -> Int) -> [a] -> [a] -> IO Int
distM ini f a b = do
        let la = length a
        let lb = length b

        arr <- newListArray ((0,0),(la,lb)) [ini i j | i<- [0..la], j<-[0..lb]] :: IO (IOArray (Int,Int) Int)

-- all on one line
        mapM_ (\(i,j) -> readArray arr (i-1,j-1) >>= \ld -> readArray arr (i-1,j) >>= \l -> readArray arr (i,j-1) >>= \d-> writeArray arr (i,j) $ minimum [l+1,d+1, ld + (f (a !! (i-1) ) (b !! (j-1))) ] ) [(i,j)| i<-[1..la], j<-[1..lb]]

        readArray arr (la,lb)