文字列間のレーベンシュタイン距離を求める(2)Haskell版

一昨日のお題を、Haskell でやってみた。

module Main where

import System.Environment (getArgs)

mValue :: [[a]] -> Int -> Int -> a
mValue m i j = (m !! i) !! j

mReplace :: [[a]] -> Int -> Int -> a -> [[a]]
mReplace m i j v = bi ++ [bj ++ [v] ++ aj] ++ ai
  where
    bi = take i m
    ji = m !! i
    ai = drop (i + 1) m
    bj = take j ji
    aj = drop (j + 1) ji

ldInit :: String -> String -> [[Int]]
ldInit s1 s2 = [[0 | j <- [0..n]] | i <- [0..m]]
  where
    m = length s1
    n = length s2

ldBuild :: ([[Int]], String, String) -> (Int, Int) -> ([[Int]], String, String)
ldBuild (m, s1, s2) (i, j) = (mReplace m i j v, s1, s2)
  where
    v = ldCalc m s1 s2 i j

ldCalc :: [[Int]] -> String -> String -> Int -> Int -> Int
ldCalc m s1 s2 i j =
  if i == 0 && j == 0 then
    0
  else if i == 0 then
    j
  else if j == 0 then
    i
  else
    let x = if (s1 !! (i-1)) == (s2 !! (j-1)) then 0 else 1
  in
  minimum [(mValue m i (j-1)) + 1, (mValue m (i-1) j) + 1, (mValue m (i-1) (j-1)) + x]

levenshteinDistance :: String -> String -> Int
levenshteinDistance s1 s2 = mValue ldM (length s1) (length s2)
  where
    m = ldInit s1 s2
    ij = [(i, j) | i <- [0..(length s1)], j <- [0..(length s2)]]
    (ldM, _, _) = foldl ldBuild (m, s1, s2) ij

main :: IO ()
main = do
  args <- getArgs
  let s1 = args !! 0
  let s2 = args !! 1
  print $ levenshteinDistance s1 s2

0 で初期化した二次元リスト(マトリックス)から foldl を使ってレーベンシュタイン距離のマトリックスを作る、というアイデアは出たんだけど、マトリックスを 0 で初期化する(ldInit)のと要素を置き換えたマトリックスを作る(mReplace)のがなかなか難産だった。結果として ldInit はリスト内包表記を使って簡潔にかけたけど、mReplace のほうはなんか泥臭くなってしまった。もっとエレガントにいかないもんだろうか。 実行例:

takatoh@nightschool $ runhaskell ld.hs apple play
4
takatoh@nightschool $ runhaskell ld.hs perl pearl
1