Find a n x m binary matrix with exactly a 1's in all rows and b 1's in all columns [Base on https://codeforces.com/contest/1360/problem/G]

So my logic is to process each column at a time and find b rows such that the total 1s in them is less than a and if there are many candidates I take the rows with least 1s.

import           Data.List
import           Data.Array
import           Control.Monad
-- import           Debug.Trace

-- debug = flip trace


main :: IO ()
main = do
    t &lt;- read &lt;$&gt; getLine -- testcases
    replicateM_ t $ do
        n : m : a : b : _ &lt;- (map read . words &lt;$&gt; getLine) :: IO [Int]
        mapM putStrLn $ case solve n m a b of
            Just mat -&gt;
                "YES"
                    : [ concatMap show [ mat ! (i, j) | j &lt;- [1 .. m] ]
                      | i &lt;- [1 .. n]
                      ]
            Nothing -&gt; ["NO"]


-- solve :: Int -&gt; Int -&gt; Int -&gt; Int -&gt; Maybe (Array (Int, Int) Int)
solve n m a b = go
    (array (1, n) . zip [1 .. n] $ repeat 0) -- rowcnts
    (array ((1, 1), (n, m)) [ ((i, j), 0) | i &lt;- [1 .. n], j &lt;- [1 .. m] ]) -- matrix
    1 -- column
  where
    go rowcnts mat c
        | length idxs /= b = Nothing -- not enough rows
        | c == m = if all (== a) (elems rowcnts') then Just mat' else Nothing -- last column to process
        | otherwise = go rowcnts' mat' (c + 1) -- recurse for next column
      where
        idxs = take b . filter (\i -&gt; (rowcnts ! i) &lt; a) $ sortOn
            (rowcnts !)
            [1 .. n] -- candidate rows
        rowcnts' = accum (+) rowcnts . zip idxs $ repeat 1 --`debug` show (elems rowcnts)
        mat'     = mat // [ ((idx, c), 1) | idx &lt;- idxs ]


I wish to ask how can I write better code which might be concise and cleaner.
You don't need to collect the rows in a matrix because you only touch each row once, and you don't need Data.Array because you only ever process a whole row. The check at the end makes the checks in the middle functionally superfluous. mapAccumL captures the explicit recursion.

main :: IO ()
main = do
  t &lt;- readLn
  replicateM_ t $ do
    [n,m,a,b] &lt;- map read . words &lt;$&gt; getLine
    mapM putStrLn $ case solve n m a of
      (colcnts, rows) | all (==b) colcnts -&gt; "YES" : map (concatMap show) rows
      _ -&gt; ["NO"]

solve :: Int -&gt; Int -&gt; Int -&gt; ([Int], [[Int]])
solve n m a = mapAccumL go (replicate m 0) [1..n] where
  go colcnts c = (zipWith (+) row colcnts, row) where row
    = map fst $ sortOn snd $ zip (replicate a 1 ++ repeat 0)
    $ map fst $ sortOn snd $ zip [1..] colcnts


Note that your entire approach doesn't obviously find a solution if there is one. Can you prove that?

Edit: I thought you process each row at a time. Easy to mix up since the approaches are symmetric :). I choose the row-at-a-time approach, since they can be directly printed.