rebeccaskinner · March 20, 2021 01:07
diff --git a/Buckets.hs b/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
	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