fork(1) download
  1. import Data.List (sortBy)
  2. import Data.Time (getCurrentTime, diffUTCTime)
  3. import Data.Function (on)
  4.  
  5. data NodeList = NodeList
  6. { nlRoot :: [Int]
  7. , nlLength :: Int
  8. , nlSum :: Int }
  9. deriving (Show)
  10.  
  11. map1 _ q [] = q
  12. map1 f q (x:xs) = f x : map1 f q xs
  13.  
  14. makeSums :: [Int] -> Int -> [NodeList] -> [NodeList]
  15. makeSums [] _ acc = acc
  16. makeSums (x:xs) maxlen acc = makeSums xs maxlen $ addNodeList $ map1 addNode acc $ filter lessLen acc
  17. where addNodeList y = (NodeList [x] 1 x) : y
  18. addNode (NodeList root len sum) = NodeList (x:root) (len+1) (sum+x)
  19. lessLen y = nlLength y < maxlen
  20.  
  21. makePairs :: [NodeList] -> [([Int], [Int])] -> [([Int], [Int])]
  22. makePairs [] acc = acc
  23. makePairs (x:xs) acc = makePairs xs $ (++) acc $ map pairNodes $ takeWhile sameSum xs
  24. where sameSum y = nlSum x == (nlSum y)
  25. pairNodes y = (nlRoot x, nlRoot y)
  26.  
  27.  
  28. solve :: [Int] -> Int -> [([Int], [Int])]
  29. solve nums maxlen = makePairs (sortBy compareNodes (makeSums nums maxlen [])) []
  30. where compareNodes x y = compare (nlSum x) (nlSum y)
  31.  
  32. makePairs1 [] acc = acc
  33. makePairs1 (x:xs) acc = makePairs1 xs $ map1 pairNodes acc $ takeWhile sameSum xs
  34. where sameSum y = nlSum x == (nlSum y)
  35. pairNodes y = (nlRoot x, nlRoot y)
  36.  
  37. solve1 nums maxlen = do
  38. let a1 = makeSums nums maxlen []
  39. let a2 = sortBy (compare `on` nlSum) a1
  40. let a3 = makePairs1 a2 []
  41. a3
  42.  
  43. bench f nums maxlen title = do
  44. begT <- getCurrentTime
  45. print $ take 1 $ reverse $ f nums maxlen
  46. endT <- getCurrentTime
  47. print $ (++) "solve1 " $ show (diffUTCTime endT begT)
  48.  
  49. test m = bench solve1 [1..m] (m `div` 2) ""
  50.  
  51. main = test 10 --print $ solve1 [1..4] 3
  52.  
Success #stdin #stdout 0s 4592KB
stdin
Standard input is empty
stdout
[([3],[2,1])]
"solve1   0.002583s"