module Matrix (saddlePoints) where import Data.Array (Ix, Array, assocs, bounds) import Data.Maybe (maybeToList) import Control.Monad (guard) import qualified Data.Set as S import qualified Data.Map as M data Axis a b = Row !a | Col !b deriving (Show, Eq, Ord) class Extrema a where ecmp :: Ord b => a -> b -> b -> Ordering data Min = Min deriving (Show) data Max = Max deriving (Show) instance Extrema Min where ecmp _ = compare instance Extrema Max where ecmp _ = flip compare data Extremes a b c = Extremes !a !b !c deriving (Show) instance (Extrema a, Ord b, Semigroup c) => Semigroup (Extremes a b c) where a@(Extremes ea xa as) <> b@(Extremes _ xb bs) = case ecmp ea xa xb of LT -> a EQ -> Extremes ea xa (as <> bs) GT -> b extrema :: (Ix a, Ix b, Ord c) => Array (a, b) c -> M.Map (Axis a b) ( Extremes Min c (S.Set (a, b)) , Extremes Max c (S.Set (a, b))) extrema = M.fromListWith (<>) . expand . assocs where expand xs = do (ix@(row, col), x) <- xs let s = S.singleton ix ax <- [Row row, Col col] return (ax, (Extremes Min x s, Extremes Max x s)) -- This is so generic with regard to extrema as it was designed -- to find either type of saddle point. The readme for this exercise -- at the time had a definition that was inconsistent with the tests. saddlePoints :: Array (Int, Int) Int -> [(Int, Int)] saddlePoints arr = do let getMin (Extremes Min a as, _) = (a, as) getMax (_, Extremes Max b bs) = (b, bs) fetch f k = maybeToList $ f <$> M.lookup k m ((r0, _c0), (r1, _c1)) = bounds arr m = extrema arr r <- [r0..r1] (rx, rixs) <- fetch getMax (Row r) rix@(_, c) <- S.toList rixs (cx, cixs) <- fetch getMin (Col c) guard $ rx == cx && rix `S.member` cixs return rix