[Haskell-cafe] Open Kattis Problem Srednji: Hints to improve my algorithm

[Dominik Bollmann]

Hi Haskell-Cafe,

I've been trying to solve the Open Kattis Problem called Srednji
recently, unfortunately without success.

Given a sequence A and a number B within that sequence this problem asks
to find all odd sub-sequences of A that, when sorted, have B as their
median in the middle. That is, from A we may remove some prefix and/or
suffix and if the resulting sub-sequence -- when sorted -- contains B in
the middle, then this sub-sequence is a solution. The problem asks to
find the number of all solutions. Check out
https://open.kattis.com/problems/srednji for the details.

My Haskell solution below tries to find the number of odd sub-sequences
by first locating the median and then repeatedly moving left and right
from that median to find larger and larger sub-sequence candidates. Each
found candidate is checked to have B in the middle when sorted in order
to become a solution. Moreover, I also extend each such candidate
further to the left (and to the right, respectively) to determine
whether these leftward or rightward extensions are solutions, too.

I think with this approach I systematically enumerate all solutions.
Unfortunately, though, this approach is too slow and times out on the
11th hidden test cases.

I'd therefore be thankful for hints about different approaches to
solving this problem more efficiently.


Thanks!

Dominik.

====================================================================

My current, slow Haskell code is this:

import Data.Maybe
import Data.Sequence (Seq, (<|), (|>))
import qualified Data.Sequence as Seq
import qualified Data.Vector.Unboxed as Vec

data SubSeq = SubSeq
  { getBalance :: {-# UNPACK #-} !Int
  , getSubSeq  :: Seq Int
  , from       :: {-# UNPACK #-} !Int
  , to         :: {-# UNPACK #-} !Int
  }

balancedSubSeqs :: [Int] -> Int -> [SubSeq]
balancedSubSeqs seq med = do
  candidate <- leftRight [val0]
  let lefts  = leftLeft candidate []
      rights = rightRight candidate []
  candidate ?: lefts ++ rights
  where
    medidx = fromJust (Vec.findIndex (== med) arr)
    val0   = SubSeq 0 (Seq.singleton med) medidx medidx
    arr    = Vec.fromList seq

    leftRight cands@(SubSeq balance seq i j : _)
      | i-1 < 0 || j+1 >= Vec.length arr = cands
      | otherwise =
        let v1       = arr Vec.! (i-1)
            v2       = arr Vec.! (j+1)
            balance' = newBalance balance v1 v2
            seq'     = (v1 <| seq) |> v2
        in leftRight (SubSeq balance' seq' (i-1) (j+1) : cands)

    leftLeft cand@(SubSeq balance seq i j) sols
      | i-2 < 0   = sols
      | otherwise =
        let v1 = arr Vec.! (i-2)
            v2 = arr Vec.! (i-1)
            balance' = newBalance balance v1 v2
            seq'     = v1 <| v2 <| seq
            newCand  = SubSeq balance' seq' (i-2) j
        in leftLeft newCand (newCand ?: sols)

    rightRight cand@(SubSeq balance seq i j) sols
      | j+2 >= Vec.length arr = sols
      | otherwise =
        let v1 = arr Vec.! (j+1)
            v2 = arr Vec.! (j+2)
            balance' = newBalance balance v1 v2
            seq'     = seq |> v1 |> v2
            newCand  = SubSeq balance' seq' i (j+2)
        in rightRight newCand (newCand ?: sols)

    newBalance old n1 n2
      | n1 < med, n2 < med = old - 2
      | n1 > med, n2 > med = old + 2
      | otherwise          = old

infixr 5 ?:
--(?:) :: SubSeq -> [SubSeq] -> [SubSeq]
x@(SubSeq b _ _ _) ?: xs
  | b == 0    = x : xs
  | otherwise = xs

main :: IO ()
main = do
  [len, median] <- fmap read . words <$> getLine
  seq <- fmap read . words <$> getLine
  let solutions = balancedSubSeqs seq median
  print (length solutions)