Edit-Distance Matrix For Levenshtein Distance Calculation

edit-distance.hs

module Fun where

import Control.Arrow ((<<<), (***))
import Data.Matrix (getElem, mapCol, mapRow, Matrix, setElem, zero)
import Data.List (foldl')

-- | initialize the edit-distance matrix.
-- λ> initEditMatrix "sitting" "kitten"
-- ( 0 1 2 3 4 5 6 )
-- ( 1 0 0 0 0 0 0 )
-- ( 2 0 0 0 0 0 0 )
-- ( 3 0 0 0 0 0 0 )
-- ( 4 0 0 0 0 0 0 )
-- ( 5 0 0 0 0 0 0 )
-- ( 6 0 0 0 0 0 0 )
-- ( 7 0 0 0 0 0 0 )
initEditMatrix :: [a] -> [a] -> Matrix Int
initEditMatrix = curry ( mapCol (const <<< flip (-) 1) 1
                       . mapRow (const <<< flip (-) 1) 1
                       . uncurry zero
                       . ((+) 1 . length *** (+) 1 . length)
                       )

-- | complete the edit-distance matrix
-- λ> distance "sitting" "kitten"
-- ( 0 1 2 3 4 5 6 )
-- ( 1 1 2 3 4 5 6 )
-- ( 2 2 1 2 3 4 5 )
-- ( 3 3 2 1 2 3 4 )
-- ( 4 4 3 2 1 2 3 )
-- ( 5 5 4 3 2 2 3 )
-- ( 6 6 5 4 3 3 2 )
-- ( 7 7 6 5 4 4 3 )
distance :: (Eq a, Ord a) => [a] -> [a] -> Matrix Int
distance as bs =
    let initMatrix = initEditMatrix as bs in
        foldl' outer initMatrix (zip as [2..])
          where
            outer matrix (a, rowIndex) = foldl' inner matrix (zip bs [2..])
              where
                inner matrix' (b, colIndex) =
                    if a == b then setElem (getElem (rowIndex - 1) (colIndex - 1) matrix')
                                           (rowIndex, colIndex)
                                           matrix'
                              else setElem v (rowIndex, colIndex) matrix'
                                where
                                  v = minimum [ getElem (rowIndex - 1)  colIndex      matrix' + 1
                                              , getElem  rowIndex      (colIndex - 1) matrix' + 1
                                              , getElem (rowIndex - 1) (colIndex - 1) matrix' + 1
                                              ]