• + 0 comments

    Irritating to get it done before timeouts, but satisfying to see it solved :-)

    haskell:

    module Main where
    
    import Data.List (sort, sortBy, span)
    import qualified Data.Map  as M
    
    type Query2Index = M.Map Int Int
    
    solve :: [Int] -> [Int] -> Query2Index
    solve _ [] = M.empty
    solve [] _ = M.empty
    solve xs qs = go mInitial iInitial totalInitial xsSorted qsSorted
      where
        mInitial = M.fromList $ [(q, -1) | q <- qs] -- (-1) is the fail value
        iInitial = 0
        totalInitial = 0
        xsSorted = sortBy (flip compare) xs -- DESC
        qsSorted = M.keys mInitial -- work of "sort qs" already done (ASC)
        {-
            Walk through the xs once, lobbing off qs from the front as we go
        -}
        go m _ _ [] _ = m -- ran out of xs
        go m _ _ _ [] = m -- ran out of qs
        go m accI accT xs' qs' =
            let
                newTotal = accT + head xs' -- running total of dec sorted xs
                newIndex = accI + 1 -- index is size of subset
                (covered, remaining) = span (<= newTotal) qs'
                -- Map.union is left-biased, so we can just add to the existing map
                m' = M.union (M.fromList [(q, newIndex) | q <- covered]) m
             in
                go m' newIndex newTotal (tail xs') remaining
    
    main :: IO ()
    main = do
        sizeLine <- getLine -- not used
        dataLine <- getLine
        let xs = map read $ words dataLine :: [Int]
    
        queryLine <- getLine -- not used
        queryLines <- getContents
        let qs = map read $ lines queryLines :: [Int]
    
        -- return a map of queries to their indecies (or -1 if not covered)
        let m = solve xs qs
    
        -- required display format, in the original order of the queries
        let result = map (m M.!) qs
        mapM_ print result