A quick explanation of how to approach a leetcode problem some folks asked about

Buckets.hs

module Main where
import Data.List (transpose)
import qualified Data.Vector as Vec
import Data.Matrix hiding (transpose)

-- Start with some sample starting bucket to test our program. It
-- contains a set of buckets each of which will have some number of
-- balls in them. The example here has 0 or 1 balls in each bucket,
-- but we could theoretically have any number >= 0.
sampleStartingBucket :: [Int]
sampleStartingBucket = [0,0,1,0,1,1]

-- Each step we're allowed to move a single ball by a single
-- unit. This means that we can calculate the number of moves required
-- to empty a given bucket, and move it's balls to the target bucket,
-- by multiplying the number of balls in the bucket by the distance to
-- the target bucket.
minimumMovesTo :: Int -> Int -> Int -> Int
minimumMovesTo currentIdx currentBallCount targetIdx  =
  let
    distance = abs (currentIdx - targetIdx)
  in currentBallCount * distance

-- Rather than a single target bucket, we would like to find the total
-- number of moves required to empty the bucket into any given target
-- bucket within some range. We can do this quite simply by computing
-- the minimum number of moves to empty the bucket into each target
-- bucket in the given range.
minimumMovesInRange :: (Int,Int) -> Int -> Int -> [Int]
minimumMovesInRange (startIdx, endIdx) currentIdx currentBallCount =
  [minimumMovesTo currentIdx currentBallCount idx | idx <- [startIdx..endIdx]]

-- For an input, which is a list of buckets containing some arbitrary
-- number of starting balls, we can calculate the
minimumMovesForEachBucket :: [Int] -> [[Int]]
minimumMovesForEachBucket buckets =
  let
    bucketsWithIndex = zip [0..] buckets
    maxIdx = (length bucketsWithIndex) - 1
  in [minimumMovesInRange (0, maxIdx) idx ballCount | (idx, ballCount) <- bucketsWithIndex]

sumMovesTo :: [[Int]] -> [Int]
sumMovesTo = map sum . transpose

-- After looking our previous implementation, you might see that we
-- have essentially re-implemented matrix multiplication. We can
-- reduce this down to a single operation by converting our original
-- input list into a row vector, and then multiplying it by matrix
-- given by the distance of each bucket from the current point of
-- origin.
withVectors :: [Int] -> Matrix Int
withVectors buckets =
  let
    v = Vec.fromList buckets
    l = Vec.length v
    r = rowVector v
    distanceMatrix = matrix l l $ \(row, col) -> abs (col - row)
  in r * distanceMatrix

main :: IO ()
main = do
  print sampleStartingBucket
  print . sumMovesTo . minimumMovesForEachBucket $ sampleStartingBucket
  putStrLn . prettyMatrix . withVectors $ sampleStartingBucket