Iterating over prefixes of multiple arrays/lists in Haskell efficiently

user_10167011 0
Viewed 251

Given three sorted arrays, I need to pick exactly m elements from them in total with the minimum sum. An additional constraint is that at least r of them must come from the first two. Here is my attempt:

solve :: [Int]->[Int]->[Int]->Int->Int->Int
solve xs ys zs m r = mini [ (xs'!!i) + (ys'!!j) + (zs'!!k) |
  i <- [max' (len ys - m) .. len xs],
  j <- [max' (r-i) .. len ys],
  k <- [max' (m-i-j) .. len zs],
  i+j+k == m,  
  i+j>=r
  ] where
      mini ws = if null ws then (-1) else minimum ws
      [xs',ys',zs'] = map (scanl (+) 0) [xs, ys, zs]
      max' x = max x 0
      len ws = min (length ws) m

(Since the arrays are sorted, I'm considering their prefixes of length i,j,k that satisfy the constraints, and summing them up using scanl.)

My worry is whether this is too inefficient, because of !!. Also, I am wondering if scanl is helping at all. Is this concern valid, and if so how can we fix it?

iteration list sum arrays haskell
2 Answers

The fundamental operation is that of merging sorted lists.

merge :: Ord a => [a] -> [a] -> [a]
merge [] ys = ys
merge xs [] = xs
merge xs@(x:xs') ys@(y:ys')
  | x <= y = x : merge xs' ys
  | otherwise = y : merge xs ys'

We can merge the first two, and combine the result with the third.

solve :: (Ord a, Num a) => [a] -> [a] -> [a] -> Int -> Int -> a
solve xs ys zs m r = sum . take m  $ prefix ++ rest
  where
    (prefix, suffix) = splitAt r (merge xs ys)
    rest = merge suffix zs
user_14776670

You could avoid !! while keeping the algorithm the same using zip:

solve xs ys zs m r = mini [ x + y + z |
    (i, x) <- zip [max' (len ys - m) .. len xs] xs',
    -- ...

I guess a naive compiler might recompute the second and third zips each time through the loops. You could defend against this source of inefficiency by binding them to names in the where block.

-- ...
    (i, x) <- indexedXs
-- ...
where
    indexedXs = zip [max' (len ys - m) .. len xs] xs'
    -- ...

(You may need to drop some elements from indexedYs and indexedZs in the later parts of the comprehension.)

Probably the best performance boost will come from making the conditions at the end of your list comprehension true by construction. The second one (i+j >= r) is already so, because you are choosing j to be at least r-i. So you can just delete that line to improve performance without changing what result you get. You can cause the first one, i+j+k==m, to be true by setting k=m-i-j (i.e. doing no iteration at all at this level). So, the end of the comprehension will look like this, instead:

    -- ...
    z <- take 1 (drop (m-i-j) zs')
    -- N.B. no conditions here
    ] where
    -- ...

Finally, I suspect you will get a bit of a win by computing the lengths of xs, ys, and zs just once instead of every time through the loops. So:

    -- use lenXs instead of len xs, etc. in the comprehension...
    ] where
    lenXs = length (take m xs)
    lenYs = length (take m ys)
    -- ...

I have also embedded a subtle change above: length (take m xs) will be slightly faster than min m (length xs) when the length of xs is much bigger than m (because it will look at much less of the list). This also gracefully deals with infinite lists, unlike your version.

Also, beware: I have thought about but not actually run the code suggestions above. Test the !$*# out of them!

user_7916040
Related