“Sequence full of colors” challenge on HackerRank

Question

This is the challenge:

You are given a sequence of N balls in 4 colors: red, green, yellow and blue. The sequence is full of colors if and only if all of the following conditions are true:

• There are as many red balls as green balls.
• There are as many yellow balls as blue balls.
• Difference between the number of red balls and green balls in every prefix of the sequence is at most 1.
• Difference between the number of yellow balls and blue balls in every prefix of the sequence is at most 1.
• Your task is to write a program, which for a given sequence prints True if it is full of colors, otherwise it prints False.

and this is my solution:

import Control.Monad (replicateM, foldM)

main :: IO ()
main = do
  n <- getLine
  list <- replicateM (read n :: Int) getLine
  mapM_ (print . isgood) list

isgood :: String -> Bool
isgood = isValidFinal . foldM step (0,0)
  where
    step (x,y) c
      | c == 'R' = wrapValidInJust (x - 1, y)
      | c == 'G' = wrapValidInJust (x + 1, y)
      | c == 'Y' = wrapValidInJust (x, y - 1)
      | c == 'B' = wrapValidInJust (x, y + 1)
    isValid (x,y) = abs x <= 1 && abs y <= 1
    wrapValidInJust xy = if isValid xy then Just xy else Nothing
    isValidFinal (Just x) = x == (0,0)
    isValidFinal Nothing = False

I'm kind of happy with main and isgood, but the step function seems very hard to read.

Any feedback?

Answer 1

• step does two things: (1) step to the next state, and (2) check that the state is valid. Its name suggests only the first meaning. The repetition of wrapValidInJust suggests that they can be separated.

• Instead of == you can use pattern-matching.

isgood :: String -> Bool
isgood = isValidFinal . foldM (check . step) (0, 0)
  where

    step (x, y) 'R' = (x-1, y)
    step (x, y) 'G' = (x+1, y)
    step (x, y) 'Y' = (x, y-1)
    step (x, y) 'B' = (x, y+1)

    check (x, y) = if abs x <= 1 && abs y <= 1 then Just (x, y) else Nothing

    ...

Another reason step is hard to read is that you need to keep track of differences, which complexifies the meaning of x and y. You can just as easily count the four colors independently and compute the difference only in check. This is also the opportunity to use a record or a map to make the update syntax more uniform:

data RGBY = RGBY { r, g, b, y :: Int }

isgood :: String -> Bool
isgood = isValidFinal . foldM (check . step) (RGBY 0 0 0 0)
  where

    step x 'R' = x { r = r x + 1 }
    step x 'G' = x { g = g x + 1 }
    step x 'Y' = x { y = y x + 1 }
    step x 'B' = x { b = b x + 1 }

    check x = if abs (r x - g x) <= 1 && abs (y x - b x) <= 1
      then Just (x, y) else Nothing

    ...