Ciao, ho provato a risolvere CSES - Shortest Routes I ma putroppo mi va in TLE su quasi tutti i testcase. Credo sia un problema di fattore costante, perché sono sicura che la complessità sia O(m\log m) e in locale ci mette circa 5 secondi. Qualcuno potrebbe aiutarmi?
import Data.Functor ((<&>))
import Data.Map qualified as Map
import Data.Maybe (fromJust, isNothing, maybeToList)
import Data.Set qualified as Set
data BinomialTree a = BinomialTree {root :: a, children :: [BinomialTree a]} deriving (Show)
newtype BinomialHeap a = BinomialHeap [Maybe (BinomialTree a)] deriving (Show)
safeSplit :: [Maybe a] -> (Maybe a, [Maybe a])
safeSplit [] = (Nothing, [])
safeSplit (x : xs) = (x, xs)
mergeTree :: (Ord a) => BinomialTree a -> BinomialTree a -> BinomialTree a
mergeTree a b
| root a <= root b = BinomialTree (root a) (b : children a)
| otherwise = BinomialTree (root b) (a : children b)
mergeHeaps :: (Ord a) => BinomialHeap a -> BinomialHeap a -> BinomialHeap a
mergeHeaps (BinomialHeap a) (BinomialHeap b) = BinomialHeap (dropWhile isNothing . reverse $ mergeLists Nothing (reverse a) (reverse b))
where
mergeLists Nothing [] [] = []
mergeLists carry [] [] = [carry]
mergeLists carry a b =
let (headA, tailA) = safeSplit a
(headB, tailB) = safeSplit b
(x, y, z) = reorder carry headA headB
(digit, newCarry)
| isNothing y = (x, Nothing)
| otherwise = (z, Just $ mergeTree (fromJust x) (fromJust y))
in digit : mergeLists newCarry tailA tailB
reorder Nothing Nothing Nothing = (Nothing, Nothing, Nothing)
reorder (Just x) Nothing Nothing = (Just x, Nothing, Nothing)
reorder Nothing (Just x) Nothing = (Just x, Nothing, Nothing)
reorder Nothing Nothing (Just x) = (Just x, Nothing, Nothing)
reorder (Just x) (Just y) Nothing = (Just x, Just y, Nothing)
reorder (Just x) Nothing (Just y) = (Just x, Just y, Nothing)
reorder Nothing (Just x) (Just y) = (Just x, Just y, Nothing)
reorder (Just x) (Just y) (Just z) = (Just x, Just y, Just z)
emptyHeap :: BinomialHeap a
emptyHeap = BinomialHeap []
isEmptyHeap :: BinomialHeap a -> Bool
isEmptyHeap (BinomialHeap x) = null x
singletonHeap :: a -> BinomialHeap a
singletonHeap x = BinomialHeap [Just (BinomialTree x [])]
pushHeap :: (Ord a) => a -> BinomialHeap a -> BinomialHeap a
pushHeap x = mergeHeaps (singletonHeap x)
popHeap :: (Ord a) => BinomialHeap a -> (Maybe a, BinomialHeap a)
popHeap heap
| isEmptyHeap heap = (Nothing, heap)
| otherwise = (Just heapMin, heapNew)
where
trees = (\(BinomialHeap trees) -> trees) heap
minIdx = getMinIdx trees 0 Nothing Nothing
getMinIdx [] _ _ minIdx = minIdx
getMinIdx (x : xs) i min minIdx
| isNothing min = getMinIdx xs (i + 1) (x >>= Just . root) (Just i)
| isNothing x || root (fromJust x) >= fromJust min = getMinIdx xs (i + 1) min minIdx
| otherwise = getMinIdx xs (i + 1) (x >>= Just . root) (Just i)
fst = take (fromJust minIdx) trees
snd = trees !! fromJust minIdx
trd = drop (1 + fromJust minIdx) trees
heapA = BinomialHeap (dropWhile isNothing (fst <> [Nothing] <> trd))
heapB = BinomialHeap (map Just (children $ fromJust snd))
heapNew = mergeHeaps heapA heapB
heapMin = root $ fromJust snd
dijkstra :: Map.Map Int [(Int, Int)] -> Set.Set Int -> BinomialHeap (Int, Int) -> Map.Map Int Int -> Map.Map Int Int
dijkstra adj vis q ans
| isEmptyHeap q = ans
| Set.member top vis = dijkstra adj vis rest ans
| otherwise = dijkstra adj (Set.insert top vis) (foldr (\(x, w) q -> pushHeap (dst + w, x) q) rest $ Map.findWithDefault [] top adj) (Map.insert top dst ans)
where
(Just (dst, top), rest) = popHeap q
solve :: Map.Map Int [(Int, Int)] -> Map.Map Int Int
solve adj = dijkstra adj Set.empty (singletonHeap (0, 0)) Map.empty
getAdj :: Int -> IO (Map.Map Int [(Int, Int)])
getAdj 0 = return Map.empty
getAdj i = getAdj (i - 1) >>= (\adj -> getLine <&> insert adj . map read . words)
where
insert adj [a, b, c] = case Map.lookup (a - 1) adj of
Nothing -> Map.insert (a - 1) [(b - 1, c)] adj
Just list -> Map.insert (a - 1) ((b - 1, c) : list) adj
main :: IO ()
main = getLine >>= (\[n, m] -> getAdj m >>= foldl (\y x -> y >> putStr (show x <> " ")) (return ()) . solve) . map read . words
