Minimum Sum of All Travel Times

module Main where
import Data.List
import Data.Ord

-- For testing purpose
import Test.QuickCheck
import Test.QuickCheck.Gen
import System.Random

type Z  = Integer
type Pt = (Z,Z)
type Dir = Pt
type Cost = Z

sumV :: [Pt] -> Pt
sumV = foldr addV (0,0)

addV :: Pt -> Pt -> Pt
addV (x,y) (u,v) = (x+u,y+v)

divV :: Pt -> Z -> Pt
divV (x,y) d = (x `div` d,y `div` d)

second :: (a,b,c) -> b
second (_,y,_) = y

third :: (a,b,c) -> c
third (_,_,z) = z

mean :: [Pt] -> Pt
mean xs = sumV xs `divV` genericLength xs

dir0 :: [Dir]
dir0 = [(-1,1) ,(0,1) ,(1,1)
       ,(-1,0)        ,(1,0)
       ,(-1,-1),(0,-1),(1,-1)]

-- Tweaking this would result 2x speed-up
restrict :: Dir -> [Dir] -> [Dir]
restrict d ds = ds

d :: Pt -> Pt -> Z
d (x,y) (u,v) = max (abs (x-u)) (abs (y-v))

costAt :: [Pt] -> Pt -> Cost
costAt xs p = sum $ map (d p) xs

cand :: Pt -> [Dir] -> [Pt]
cand p ds = map (addV p) ds

step :: [Pt] -> (Pt,Cost,[Dir]) -> (Pt,Cost,[Dir])
step xs (p,c,ds)
    |  c' < c    =  (p',c', restrict d ds)
    |  otherwise =  (p ,c , [])
    where
      (p',c',d) = minimumBy (comparing second) $ map (look xs p) ds

look :: [Pt] -> Pt -> Dir -> (Pt,Cost,Dir)
look xs p d = (p',costAt xs p',d)
    where p' = p `addV` d

-- seems like a reasonable guess
start :: [Pt] -> (Pt,Cost,[Dir])
start xs = (p0,c0,dir0)
    where
      p0  = mean xs
      c0  = costAt xs p0


opt :: [Pt] -> (Pt,Cost,[Dir])
opt xs =  head $ dropWhile (not . null . third) $ iterate (step xs) $ start xs

optDebug :: [Pt] -> (Bool,[(Pt,Cost,[Dir])])
optDebug xs =  (check xs p,head ns : ms)
    where
      (ms,ns) = span (not . null . third) $ iterate (step xs) $ start xs
      (p,c,_) = head ns
-------------------------------------------------
-- test code
-------------------------------------------------

check :: [Pt] -> Pt -> Bool
check xs p = costAt xs p <= minimum (map (costAt xs) (cand p dir0))

prop_opt :: [Pt] -> Bool
prop_opt xs = let (p,c,_) = opt xs
              in check xs p

rollDice :: IO Int
rollDice = getStdRandom random

randomPts :: Int -> IO [Pt]
randomPts n = do rollDice
                 gen <- getStdGen
                 return $ unGen (vectorOf n arbitrary) gen 1000

main = do xs <- randomPts 100000
          print $ optDebug xs

sample0,sample1,sample2 :: [Pt]
sample0 = [(0,0)]
sample1 = [(0,0),(1,1)]
sample2 = [(0,2),(1,2),(2,2)
          ,      (1,1)
          ,            (2,0) ]

sample3 = [(0,3)
          ,(0,2),(1,2)
          ,(0,1),(1,1),(2,1)
          ,(0,0),(1,0),(2,0),(3,0)]


sample4 = [(0,3)
          ,(0,2),(1,2)
          ,(0,1),(1,1),(2,1)
          ,(0,0),(1,0),(2,0),(3,0)]

sample5 = [ (-3, 0), (-2, 0), (-1, 0), (0, 0), (10, 0), (50, 0) ]