これはちょっとフェアじゃないんではないか。
HaskellのやつをPythonのっぽく普通に書いたらこうなった。
import Control.Applicative hiding (empty)
import qualified Data.HashMap.Strict as HM
import Data.Maybe
data Trie v
= Node (Maybe v) (HM.HashMap Char (Trie v))
empty :: Trie v
empty = Node Nothing HM.empty
insert :: Trie v -> String -> v -> Trie v
insert (Node _ hm) "" v = Node (Just v) hm
insert (Node m hm) (c:cs) v =
let t = HM.lookupDefault empty c hm in
Node m (HM.insert c (insert t cs v) hm)
find :: Trie v -> String -> Maybe (Trie v)
find t "" = Just t
find (Node _ hm) (c:cs) = do
t <- HM.lookup c hm
find t cs
allPrefixes :: Trie v -> [v]
allPrefixes (Node v hm) =
maybeToList v ++ (concatMap allPrefixes $ HM.elems hm)
autocomplete :: Trie v -> String -> [v]
autocomplete t word =
fromMaybe [] $ allPrefixes <$> find t word
main :: IO ()
main = do
ws <- lines <$> readFile "/usr/share/dict/words"
let trie = foldl (\t s -> insert t s s) empty ws
print $ autocomplete trie "inn"