二分探索は Brent Yorgey 先生に従う

Competitive programming in Haskell: better binary search

競プロ典型 90 問 51 日目 - Typical Shop（★5）

import Data.Array.Unboxed
import qualified Data.ByteString.Char8 as C
import Data.List

main = sol <$> ints <*> ints >>= print

ints = unfoldr (C.readInt . C.dropWhile (==' ')) <$> C.getLine

sol [n,k,p] as = f n k p b c 1
  where
  bs = sort as
  cs = scanl' (+) 0 bs
  b = listArray (1,n) bs :: UArray Int Int
  c = listArray (0,n) cs :: UArray Int Int

f n k p b c i
  | k==1                = (-) (bs b p i n) i
  | c!(i-1+k)-c!(i-1)>p = 0
  | c!n-c!(n-k)<=p      = comb (n-i+1) k
  | otherwise           = sum $ map (\j -> f n (k-1) (p-b!j) b c (j+1)) [i..n-k+1]

bs b p l r
  | b!l>p     = l
  | b!r<=p    = r+1
  | otherwise = snd $ search binary ((>p) . (b!)) l r

search mid p = go
  where
  go l r = case mid l r of
    Nothing -> (l,r)
    Just m
      | p m       -> go l m
      | otherwise -> go m r

binary l r
  | r-l > 1   = Just ((l+r) `div` 2)
  | otherwise = Nothing

comb p q
  | q==0      = 1
  | p-q<q     = comb p (p-q)
  | otherwise = comb p (q-1)*(p+1-q) `div` q