import Data.List import Data.Foldable data Bundle = Bundle { count :: Int, col :: Int } deriving (Show) toLeft :: Bundle -> Bundle toLeft (Bundle c i) = Bundle c (i - 1) toRight :: Bundle -> Bundle toRight (Bundle c i) = Bundle c (i + 1) overlap :: Bundle -> Bundle -> Bundle overlap (Bundle c1 i) (Bundle c2 _) = Bundle (c1 + c2) i instance Eq Bundle where (==) b1 b2 = col b1 == col b2 instance Ord Bundle where b1 `compare` b2 = col b1 `compare` col b2 mergeBundles :: Bundle -> Bundle -> (Bundle, Maybe Bundle) mergeBundles b1 b2 | b1 < b2 = (b1, Just b2) | b1 > b2 = (b2, Just b1) | b1 == b2 = (overlap b1 b2, Nothing) mergeBundleLists :: [Bundle] -> [Bundle] -> [Bundle] mergeBundleLists b1s [] = b1s mergeBundleLists [] bs2 = bs2 mergeBundleLists (b1:b1s) (b2:b2s) = minB:mergeBundleLists b1s' b2s where (minB, maxB) = mergeBundles b1 b2 b1s' = maybe b1s (prependBundle b1s) maxB prependBundle :: [Bundle] -> Bundle -> [Bundle] prependBundle [] b = [b] prependBundle (fb:beams) b = min_b:(toList max_b ++ beams) where (min_b, max_b) = mergeBundles b fb mergeSplit :: Bundle -> ([Bundle], Int) -> ([Bundle], Int) mergeSplit b (beams, splitCount) = (toLeft b : prependBundle beams (toRight b), splitCount + 1) mergeMiss :: Bundle -> ([Bundle], Int) -> ([Bundle], Int) mergeMiss b (beams, splitCount) = (prependBundle beams b, splitCount) passBundles :: [Int] -> [Bundle] -> ([Bundle], Int) passBundles [] beams = (beams, 0) passBundles _ [] = ([], 0) passBundles (s:splitters) (b@(Bundle size i):bundleList) | s < i = passBundles splitters (b:bundleList) | s > i = mergeMiss b (passBundles (s:splitters) bundleList) | s == i = mergeSplit b (passBundles splitters bundleList) sources :: String -> [Bundle] sources = map (Bundle 1) . elemIndices 'S' splitters :: String -> [Int] splitters = elemIndices '^' passLine :: [([Bundle], Int)] -> String -> [([Bundle], Int)] passLine (beams:beamss) line = (res, splitCount + snd beams):beams:beamss where spls = splitters line bms = mergeBundleLists (fst beams) (sources line) (res, splitCount) = passBundles spls bms countBundles :: [Bundle] -> Int countBundles bs = sum (map count bs) passLines :: String -> [([Bundle], Int)] passLines inputStr = foldl passLine [([], 0)] (lines inputStr) main :: IO() main = do inputStr <- getContents print ((\(bs, c) -> (c, countBundles bs)) $ head $ passLines inputStr)