dmjio · February 7, 2024 21:04 · Feb 7, 2024
diff --git a/SuffixTree.hs b/SuffixTree.hs
@@ -0,0 +1,116 @@
+import           Data.List               (foldl', tails)
+import           Data.Map.Internal.Debug
+import qualified Data.Map.Strict         as M
+import           Data.Maybe
+import           Data.Monoid
+import qualified Data.Set as S
+
+data TrieValue v
+  = TrieValue
+  { next :: Trie v
+  , value :: v
+  } deriving (Show, Eq)
+
+instance Semigroup v => Semigroup (TrieValue v) where
+  TrieValue n1 tv1 <> TrieValue n2 tv2 =
+    TrieValue (n1 <> n2) (tv1 <> tv2)
+
+-- | suffix trie, a trie populated with suffixes
+newtype Trie v = Trie (M.Map Char (TrieValue v))
+  deriving (Eq, Show)
+
+type Todos = Trie (S.Set Int)
+
+instance Monoid k => Monoid (Trie k) where
+  mempty = empty
+
+instance Semigroup k => Semigroup (Trie k) where
+  Trie m <> Trie n = Trie (m <> n)
+
+empty :: Trie v
+empty = Trie mempty
+
+-- Assumes v is todo ID and String has no spaces, only a-z and '-'
+-- this inserts all suffixes into the trie
+fromList :: Monoid v => [(v, String)] -> Trie v
+fromList xs = foldr insertWithTails empty xs
+  where
+    insertWithTails (k,v) m =
+      foldr (Main.insert k) m (tails v)
+
+fromList_ :: String -> Trie ()
+fromList_ vs = foldr (Main.insert ()) mempty (tails vs)
+
+insert
+  :: Monoid v
+  => v
+  -> String
+  -> Trie v
+  -> Trie v
+insert _ [] t            = t
+insert v (x:xs) (Trie m) =
+  case M.lookup x m of
+    Nothing ->
+      Trie (M.insertWith (<>) x
+            (TrieValue (insert v xs mempty) v) m)
+    Just n ->
+      Trie (M.insertWith (<>) x
+            (TrieValue (insert v xs (next n)) v) m)
+
+search :: Monoid v => String -> Trie v -> v
+search [] _ = mempty
+search xs m = go xs m
+  where
+    go [x] (Trie n) =
+      case M.lookup x n of
+        Nothing -> mempty
+        Just g -> value g
+
+    go (x:xs) (Trie n) =
+      case M.lookup x n of
+        Nothing -> mempty
+        Just g -> go xs (next g)
+
+size :: Trie v -> Int
+size (Trie m) =
+  case M.size m of
+    0 -> 0
+    n -> n + sum (size . next <$> M.elems m)
+
+deleteById :: Int -> Todos -> Todos
+deleteById key (Trie m) = Trie $
+  M.fromList
+  [ (c, TrieValue (deleteById key n) (S.delete key keys))
+  | (c, TrieValue n keys) <- M.toList m
+  , S.singleton key /= keys
+  ]
+
+main :: IO ()
+main = do
+  let todos =
+        [ ( S.singleton 1
+          , "eat"
+          ),
+          ( S.singleton 2
+          , "drink"
+          ),
+          ( S.singleton 3
+          , "sleep"
+          ),
+          ( S.singleton 4
+          , "yay"
+          ),
+          ( S.singleton 5
+          , "foo"
+          )
+        ]
+  let ts = fromList todos
+  print $ search "e" ts
+  print $ search "s" ts
+  print $ search "ink" ts
+  print $ search "eat" ts
+  print $ search "a" ts
+  print $ search "at" (deleteById 4 ts)
+
+
+