Bear.hs

module Bear (main, findMinSubStr) where
import qualified Data.Map.Strict as Map
import qualified Data.List as List
import qualified Data.ByteString as BS
import Data.Maybe (fromMaybe)
import Data.ByteString.Char8 (readInt)
import GHC.Word (Word8)

type CharFreq = Map.Map Word8 Int

extraChars :: Int -> BS.ByteString -> CharFreq
extraChars geneLength s =
    let numNeeded = geneLength `div` 4
        counts = Map.fromListWith (+) $ BS.foldl' (\a c -> (c, 1) : a) [] s
        adjusted = Map.map (\x -> x - numNeeded) counts
     in adjusted

haveWhatINeed :: CharFreq -> Bool
haveWhatINeed m = List.all (<= 0) $ Map.elems m

rubberband :: Int -> Int -> Int -> Int -> CharFreq -> BS.ByteString -> Int -> Int
rubberband lengthOfGene targetLength start end extra gene minSubStrLen
  | minSubStrLen == targetLength = minSubStrLen
  | start >= lengthOfGene || end >= lengthOfGene = minSubStrLen
  | haveWhatINeed extra =
      rubberband
        lengthOfGene targetLength (start + 1) end
        (Map.adjust (+ 1) (BS.index gene start) extra)
        gene (min (end - start) minSubStrLen)
  | otherwise = rubberband
                    lengthOfGene targetLength start (end + 1)
                    (Map.adjust (\x -> x - 1) (BS.index gene end) extra)
                    gene minSubStrLen

findMinSubStr :: Int -> BS.ByteString -> Int
findMinSubStr len gene = let extra = extraChars len gene
                             targetLength = List.sum $ List.filter (> 0) $ Map.elems extra
                          in rubberband len targetLength 0 0 extra gene len

main :: IO ()
main = do
    len <- BS.getLine
    gene <- BS.getLine
    let lenInt = fst $ fromMaybe (error "bad input") (readInt len)
    print $ findMinSubStr lenInt gene