코딩 인터뷰 완전 분석 215쪽 문제 18.10 풀이

search.hs

import System.IO (hFlush, stdout)
import Data.List (intercalate)
import qualified Data.Set as S
import qualified Data.Map as M

type Word = String
type Dict = [Word]

-- 특정 글자를 기준으로 스트링을 잘라 주는 함수
-- splitWith d s 는 d에 포함된 글자가 나오면 그 글자를 기준으로 스트링을 자른다
-- splitWith "," "a,b,c" = ["a", "b", "c"]
-- splitWith ", " "a, b, c" = ["a", "b", "c"]
splitWith :: (Eq a) => [a] -> [a] -> [[a]]
splitWith d [] = []
splitWith d s  = let (e, rest) = break (`elem` d) s
                     next = dropWhile (`elem` d) rest
                  in e : splitWith d next

-- 파일을 읽어들여 사전을 구성하는 IO Monad
loaddict :: [Char] -> IO Dict
loaddict filename = do s <- readFile filename
                       return $ splitWith "\n" s

-- 시작 단어를 한글자만 변경시킨 모든 단어들을 사전 순으로 반환합니다.
-- 사전 순으로 반환하는 이유는, 밑의 intersect_sorted 함수와 함께 사전 탐색을 보다 빨리 하기 위해서입니다.
-- 편의상 입력 단어의 길이는 같다고 가정합니다.
modify_one :: [Char] -> [[Char]]
modify_one []     = []
modify_one (x:xt) = [y:xt | y <- ['a' .. pred x]] ++ [x:yt | yt <- modify_one xt] ++ [y:xt | y <- [succ x .. 'z']]

-- 두개의 "정렬된" 리스트를 입력으로 받아 양쪽 모두에 존재하는 리스트를 출력합니다
-- intersect_sorted [3, 6, 9, 12] [1, 3, 5, 7, 9] = [3, 9]
intersect_sorted :: (Ord a) => [a] -> [a] -> [a]
intersect_sorted xs     []     = []
intersect_sorted []     ys     = []
intersect_sorted (x:xs) (y:ys) = case compare x y of
                                    LT -> intersect_sorted xs (y:ys)
                                    EQ -> x : intersect_sorted xs ys
                                    GT -> intersect_sorted (x:xs) ys

-- 두 단어 사이의 거리를 계산 합니다. 두 단어의 길이는 같다고 가정합니다.
-- distance "lamp" "lime" = 2
distance :: (Eq a) => [a] -> [a] -> Int
distance []     []     = 0
distance (x:xs) (y:ys) | x == y    = distance xs ys
                       | otherwise = 1 + distance xs ys

-- 시작 단어와 목적 단어가 있을 때, 시작 단어에서 최단거리로 목적 단어까지 가는 가장 빠른 경로를 반환
-- A* search를 수행하여, 경로를 찾으면 Just(경로)를 반환, 없으면 Nothing을 반환합니다.
shortest_path :: Dict -> Word -> Word -> Maybe [Word]
shortest_path dict start end = a_star_search (S.singleton (distance start end + 1, [end])) (S.singleton end)
  where
    a_star_search :: S.Set (Int, [Word]) -> S.Set Word -> Maybe [Word]
    a_star_search q visited
      = do ((dist, ws), q') <- S.minView q
           let w = head ws
           if w == start
              then return ws
              else let n = length ws
                       next_ws = filter (`S.notMember` visited) $ intersect_sorted dict $ modify_one w
                       cost_ws = map (\w' -> (distance start w' + n, w':ws)) next_ws
                    in a_star_search (S.fromList cost_ws `S.union` q') (S.fromDistinctAscList next_ws `S.union` visited)

-- 두 단어를 받아 경로를 출력하는 함수
-- 단어 길이가 4일 경우 word4.txt, 5일 경우 word5.txt를 이용하며 그 외에는 에러를 출력합니다.
print_shortest :: Dict -> Dict -> [Char] -> IO ()
print_shortest dict4 dict5 line = do let [start, end] = splitWith ", " line
                                     let dict = case (length start, length end) of
                                                  (4,4) -> dict4
                                                  (5,5) -> dict5
                                                  _     -> error "bad word length"
                                     case (start `elem` dict, end `elem` dict) of
                                        (False, _) -> error $ show start ++ " not in dictionary"
                                        (_, False) -> error $ show end ++ " not in dictionary"
                                        _          -> case shortest_path dict start end of
                                                        Nothing -> putStrLn $ "impossible"
                                                        Just ps -> putStrLn $ "output: " ++ intercalate " -> " ps

main :: IO ()
main = do dict4 <- loaddict "word4.txt"
          dict5 <- loaddict "word5.txt"
          putStr "input: "
          hFlush stdout
          input <- getContents
          case splitWith "\r\n" input of
            []       -> putStrLn ""
            (line:_) -> print_shortest dict4 dict5 line