Haskell code template for competitive programming

codeforces_template.hs

{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE OverloadedStrings #-}

{-|
Module      : Main
Description : A Haskell code template designed for competitive programming, targeting platforms like CodeForces.

Presenting a Haskell code template designed for competitive programming, targeting platforms like CodeForces.
The template offers a modular structure to handle diverse tasks including input/output management, parsing, test case segmentation, and formatting.
This approach facilitates the formulation of solutions through elegant and efficient pure functions, optimizing code clarity and maintainability.
Elevate your competitive programming prowess using this streamlined template.

To use this template, follow the steps below:

1. simply copy the code from [gist](https://gist.github.com/mihassan/3cf472fee6383bd4a3d62274a28e095e).
2. Then, update the @solution@ function with your solution.
3. Finally, update the @main@ function with the appropriate configuration.
4. Optionally, if existing Parsers and Formatters are not enough for the problem, you can add new Parser and Formatter for you specific problem.
-}
module Main where

import           Data.Bool
import qualified Data.ByteString.Char8         as BS
import           Data.List
import           Data.Maybe
import           System.IO

-- * Template Code

-- ** Type definitions

-- | Parser type class to handle lines of input for a single test case.
class Parser i where
  parse :: [BS.ByteString] -> i

-- | Formatter type class to format output for a single test case.
class Formatter o where
  format :: o -> BS.ByteString

-- | Splitter type to segment the input into multiple test cases.
data Splitter
  = NoSplitting -- ^ No splitting, the entire input is a single test case.
  | EqualSized -- ^ Split into equal sized test cases. Number of test cases is provided in the first line of the input.
  | FixedSized Int -- ^ Split into fixed sized test cases. Number of test cases is provided in the first line of the input.
  | DynamicSized ([Int] -> Int) -- ^ Split into dynamic sized test cases based on the first line of the test case. Number of test cases is provided in the first line of the input.

-- | Configuration to handle input/output and splitting.
data Config = Config
  { input    :: String -- ^ Input file path. If empty, stdin is used.
  , output   :: String -- ^ Output file path. If empty, stdout is used.
  , splitter :: Splitter -- ^ Splitter to segment the input into multiple test cases.
  }

-- | Default configuration with stdin/stdout and equal sized splitting based on the first line of the input.
with :: Config
with = Config { input = "", output = "", splitter = EqualSized }

-- | Solution type to handle the actual solution. It is a pure function that the user needs to provide.
type Solution i o = i -> o

-- ** Parsers

-- | Parse as lines of Strings.
instance Parser [BS.ByteString] where
  parse = id

-- | Parse as lines of words of Strings.
instance Parser [[BS.ByteString]] where
  parse = map BS.words

-- | Parse as lines of multiple Integer per lines.
instance Parser [[Integer]] where
  parse = (fmap . fmap) (unsafeRead BS.readInteger) . parse

-- | Parse the first line as String and the rest as lines of multiple Integers per lines.
-- This is useful when the first line needs to be handled differently.
instance Parser (BS.ByteString, [[Integer]]) where
  parse = fmap parse . fromJust . uncons

-- | Helper combinator to be used with ByteString read* functions.
unsafeRead :: (a -> Maybe (b, c)) -> a -> b
unsafeRead f = fst . fromJust . f

-- ** Formatters

-- | Format a String as itself.
instance Formatter BS.ByteString where
  format = id

-- | Format an Integer as String.
instance Formatter Integer where
  format = BS.pack . show

-- | Format a Bool value as capital String - "YES" and "NO".
instance Formatter Bool where
  format = bool "NO" "YES"

-- | A wrapper over a List to control how it is formatted.
data ListOf a
  = WordsOf [a] -- ^ Format as words separated by space.
  | LinesOf [a] -- ^ Format as lines separated by newline.

-- | Format a list of values either on a single line or multiple lines.
instance (Formatter a) => Formatter (ListOf a) where
  format (WordsOf xs) = BS.unwords $ format <$> xs
  format (LinesOf xs) = BS.init . BS.unlines $ format <$> xs

-- ** Splitters

-- | Split the input into multiple test cases based on the splitter provided. 
splitWith :: Splitter -> [BS.ByteString] -> [[BS.ByteString]]
splitWith NoSplitting xs       = [xs]
splitWith EqualSized  (x : xs) = splitWith (FixedSized n) (x : xs)
  where n = length xs `div` unsafeRead BS.readInt x
splitWith (FixedSized   n) (_ : xs) = dynamicSizedSpliiter (const n) xs
splitWith (DynamicSized f) (_ : xs) = dynamicSizedSpliiter f xs

-- | Split the input into multiple test cases.
-- The user needs to provide a function that computes the size of the test case based on the first line.
-- It could just be n-th Int in the first line, or a combination of multiple Ints.
dynamicSizedSpliiter :: ([Int] -> Int) -> [BS.ByteString] -> [[BS.ByteString]]
dynamicSizedSpliiter _ [] = []
dynamicSizedSpliiter f xs = fst : dynamicSizedSpliiter f rest
  where
    header      = BS.words $ head xs
    size        = f $ unsafeRead BS.readInt <$> header
    (fst, rest) = splitAt size xs

-- ** The main solve function

-- | Combining all together, we have the solve function.
--
--   1. Split into lines.
--   2. Split the lines into multiple test cases based on the splitter provided.
--   3. For each test case, parse, apply solution, and format.
--   4. Finally join the output to show.
solve :: (Parser i, Formatter o) => Config -> Solution i o -> IO ()
solve Config {..} solution = do
  i <- if null input then return stdin else openFile input ReadMode
  o <- if null output then return stdout else openFile output WriteMode
  hInteract i o $ solvePure splitter

-- | Generic interact function with the given input/output handles.
hInteract :: Handle -> Handle -> (BS.ByteString -> BS.ByteString) -> IO ()
hInteract i o f = BS.hGetContents i >>= BS.hPutStr o . f >> hFlush o

-- | Pure function to solve the problem.
solvePure :: Splitter -> BS.ByteString -> BS.ByteString
solvePure splitter =
  BS.unlines . map solveTestCase . splitWith splitter . BS.lines

-- | Solve a single test case.
solveTestCase :: [BS.ByteString] -> BS.ByteString
solveTestCase = format . solution . parse

-- * User Code

-- | Sample solution for CodeForces [problem 900A](https://codeforces.com/contest/900/problem/A).
-- Few examples for different type signatures are provided below.
--
-- @ solution :: Solution [BS.ByteString] BS.ByteString @
--
-- @ solution :: Solution [[BS.ByteString]] Integer @
--
-- @ solution :: Solution (BS.ByteString, [[Integer]]) Bool @
--
-- @ solution :: Solution [[Integer]] (WordsOf Integer) @
--
-- @ solution :: Solution [[Integer]] (LinesOf Integer) @
--
-- @ solution :: Solution [[Integer]] (LinesOf (WordsOf Integer)) @
solution :: Solution [[Integer]] Bool
solution ([n] : ps) = n <= 1 || p <= 1
  where
    xs = head <$> ps
    n  = countIf (< 0) xs
    p  = countIf (> 0) xs

-- | Count the number of elements in a list that satisfy the predicate.
countIf :: Eq a => (a -> Bool) -> [a] -> Int
countIf p = length . filter p

-- | Main function to run the solution. Update the Config as necessary.
-- Few examples are provided below.
--
-- Reads multiple test cases from stdin and writes to stdout.
-- The number of test cases is provided in the first line of the input.
-- Rest of the input is split into equal sized test cases.
-- 
-- @
--   main = solve with solution
-- @
--
-- Reads a single test case from stdin and writes to stdout.
-- 
-- @
--   main = solve with { splitter = NoSplitting } solution
-- @
--
-- Reads multiple test cases, each containing 2 lines, from a "input.txt" and writes to stdout.
--
-- @
--   main = solve with { input = "input.txt", splitter = FixedSized 2 } solution
-- @
--
-- Reads multiple test cases from stdin and writes to stdout.
-- Size of each test case is given by first Int in the first line of the test case.
--
-- @
--   main = solve with { splitter = DynamicSized (!! 0) } solution
-- @
main :: IO ()
main = solve with { splitter = NoSplitting } solution