Bear and Steady Gene

  • + 0 comments

    Haskell:

    module Main where
    
    import qualified Data.Map as M
    import Data.List (find)
    import qualified Data.Vector  as V
    
    type Counts = M.Map Char Int
    type Window = (Int, Int, Counts)
    
    symbols :: [Char]
    symbols = ['A', 'C', 'G', 'T']
    
    initCounts :: Counts
    initCounts = M.fromList [(c, 0) | c <- symbols]
    
    getCounts :: String -> Counts
    getCounts = foldl (\m c -> M.insertWith (+) c 1 m) initCounts
    
    extras :: String -> Maybe Counts
    extras [] = Nothing
    extras s =
        let
            n = length s `div` 4  -- we're asured it's a multiple of 4
            acc0 = M.fromList [(c, 0) | c <- symbols]
            counts = getCounts s
            adjusted = M.map (\x -> x - n) counts
            es = M.filter (> 0) adjusted
         in
            case length es of
                0 -> Nothing  -- already balanced
                _ -> Just es  -- work to do
    
    satisfied :: Counts -> Counts -> Bool
    satisfied required actual = all (\c -> actual M.! c >= required M.! c) (M.keys required)
    
    soln0 :: Counts -> String -> Maybe (Int, Counts)
    soln0 es s = do
        let dists = drop 1 $ scanl (\m c -> M.insertWith (+) c 1 m) initCounts s
        let indexedDists = zip [0..] dists
        find (\(j, d) -> satisfied es d) indexedDists
    
    slideUntilSuccess :: V.Vector Char -> Counts -> Window -> Maybe Window
    slideUntilSuccess v required (i, j, counts) = do
        addC <- v V.!? (j + 1)  -- return Nothing if j + 1 is out of bounds
        let cminus = M.insertWith (+) (v V.! i) (-1) counts
        let counts' = M.insertWith (+) addC 1 cminus
        (if satisfied required counts' then Just (i + 1, j + 1, counts') else slideUntilSuccess v required (i + 1, j + 1, counts'))
    
    shrinkUntilFail :: V.Vector Char -> Counts -> Window -> Window
    shrinkUntilFail v required (i, j, counts) =
        if satisfied required counts
            then let counts' = M.insertWith (+) (v V.! i) (-1) counts
                in shrinkUntilFail v required (i + 1, j, counts')
            else (i, j, counts)
    
    runner :: V.Vector Char -> Counts -> Window -> Window
    runner v required w = do
        let w' = shrinkUntilFail v required w
        let w'' = slideUntilSuccess v required w'
        case w'' of
            Nothing -> w'
            Just ws -> runner v required ws
    
    solve :: String -> Maybe Window
    solve s = do
        es <- extras s
        (j0, dist0) <- soln0 es s
        let v = V.fromList s
        return $ runner v es (0, j0, dist0)
    
    main :: IO ()
    main = do
        n <- readLn :: IO Int
        s <- getLine
        case solve s of
            Nothing -> putStrLn "0"
            Just (i, j, _) -> print (j-i+2)