Efficient graph breadth-first search?

After studying graph-related materials in Haskell, I managed to solve the graph bipartite problem on CSES. However, my solution was not efficient enough to pass all test cases.

I would appreciate any suggestions for improvement. Thank you.

Here is the problem statement: https://cses.fi/problemset/task/1668

Below is my code (stolen from "King, David Jonathan (1996) Functional programming and graph algorithms. PhD thesis"):

```hs {-# LANGUAGE RankNTypes #-}

import Debug.Trace import qualified Data.ByteString.Char8 as B import Control.Monad import Data.Array import Data.List import Data.Set qualified as Set import Data.Set (Set) import Data.Maybe

type Vertex = Int type Edge = (Vertex, Vertex) type Graph = Array Vertex [Vertex]

vertices :: Graph -> [Vertex] vertices = indices

edges :: Graph -> [Edge] edges g = [ (v, w) | v <- vertices g , w <- g!v ]

mkgraph :: (Vertex, Vertex) -> [Edge] -> Graph mkgraph bounds edges = accumArray (flip (:)) [] bounds (undirected edges) where undirected edges = concatMap ((v, w) -> [(v, w), (w, v)]) edges

data Tree a = Node a (Forest a) type Forest a = [Tree a]

generateT :: Graph -> Vertex -> Tree Vertex generateT g v = Node v (generateF g (g!v))

generateF :: Graph -> [Vertex] -> [Tree Vertex] generateF g vs = map (generateT g) vs

bfsPrune :: [Tree Vertex] -> Set Vertex -> ([Tree Vertex], Set Vertex) bfsPrune ts q = let (us, ps, r) = traverseF ts (q:ps) in (us, r) where traverseF [] ps = ([], ps, head ps) traverseF (Node x ts : us) (p:ps) | Set.member x p = traverseF us (p:ps) | otherwise = let (ts', qs, q) = traverseF ts ps (us', ps', p') = traverseF us ((Set.insert x p) : qs) in (Node x ts' : us', ps', Set.union q p')

bfs :: Graph -> [Vertex] -> Set Vertex -> ([Tree Vertex], Set Vertex) bfs g vs p = bfsPrune (generateF g vs) p

bff :: Graph -> [Vertex] -> Set Vertex -> [Tree Vertex] bff g [] p = [] bff g (v:vs) p | Set.member v p = bff g vs p | otherwise = let (ts, p') = bfs g [v] p in ts <> bff g vs p'

preorderF :: forall a. [Tree a] -> [a] preorderF ts = concatMap preorderT ts where preorderT (Node x ts) = x : preorderF ts

type Color = Int

annotateF :: forall a. Color -> [Tree a] -> [Tree (a, Color)] annotateF n ts = map (annotateT n) ts where switch n = if n == 1 then 2 else 1 annotateT n (Node x ts) = let ts' = annotateF (switch n) ts in Node (x, n) ts'

colorArr :: Graph -> Array Vertex Color colorArr g = let ts = bff g (vertices g) Set.empty in array (bounds g) (preorderF (annotateF 1 ts))

isBipartite :: Graph -> (Bool, Array Vertex Color) isBipartite g = let color = colorArr g in (and [color!v /= color!w | (v, w) <- edges g], color)

readInt :: B.ByteString -> Int readInt = fst . fromJust . B.readInt

ints :: IO (Int, Int) ints = do [x, y] <- B.words <$> B.getLine pure (readInt x, readInt y)

main :: IO () main = do (v, e) <- ints es <- replicateM e ints let g = mkgraph (1,v) es (b, color) = isBipartite g if b then do putStrLn $ unwords $ map (\v -> show $ color!v) [1..v] else putStrLn "IMPOSSIBLE" ```