112 lines
3.9 KiB
Haskell
112 lines
3.9 KiB
Haskell
module Main (main) where
|
|
|
|
import Data.List (findIndex, groupBy)
|
|
import Data.Maybe (fromJust)
|
|
|
|
-- | Part 1
|
|
data DiskItem = File Int Int | Free Int deriving (Show, Eq)
|
|
|
|
isFree :: DiskItem -> Bool
|
|
isFree (Free _) = True
|
|
isFree _ = False
|
|
|
|
isFile :: DiskItem -> Bool
|
|
isFile (File _ _) = True
|
|
isFile _ = False
|
|
|
|
fileLen :: DiskItem -> Int
|
|
fileLen (File _ len) = len
|
|
fileLen (Free len) = len
|
|
|
|
fileId :: DiskItem -> Int
|
|
fileId (File fId _) = fId
|
|
fileId (Free _) = error "Free has no id!"
|
|
|
|
readDiskMap :: String -> [DiskItem]
|
|
readDiskMap = readDiskMap' 0
|
|
where
|
|
readDiskMap' :: Int -> String -> [DiskItem]
|
|
readDiskMap' _ [] = []
|
|
readDiskMap' _ ['\n'] = []
|
|
readDiskMap' fileId (i : is)
|
|
| even fileId = File (fileId `div` 2) (read [i] :: Int) : readDiskMap' (fileId + 1) is
|
|
| otherwise = Free (read [i] :: Int) : readDiskMap' (fileId + 1) is
|
|
|
|
expandDiskMap :: [DiskItem] -> [DiskItem]
|
|
expandDiskMap [] = []
|
|
expandDiskMap ((File fileId len) : is) = replicate len (File fileId 1) ++ expandDiskMap is
|
|
expandDiskMap ((Free len) : is) = replicate len (Free 1) ++ expandDiskMap is
|
|
|
|
showDiskMap :: [DiskItem] -> IO ()
|
|
showDiskMap [] = putStrLn ""
|
|
showDiskMap ((File fileId len) : is) = putStr (concat [show fileId | _ <- [1 .. len]]) >> showDiskMap is
|
|
showDiskMap ((Free len) : is) = putStr (replicate len '.') >> showDiskMap is
|
|
|
|
defragmentStep :: [DiskItem] -> [DiskItem]
|
|
defragmentStep diskMap =
|
|
let diskSize = length diskMap
|
|
firstFree = fromJust $ findIndex isFree diskMap
|
|
lastFile = diskSize - 1 - fromJust (findIndex (not . isFree) (reverse diskMap))
|
|
in take firstFree diskMap ++ [diskMap !! lastFile] ++ take (lastFile - firstFree - 1) (drop (firstFree + 1) diskMap) ++ [Free 1] ++ take (diskSize - lastFile) (drop (lastFile + 1) diskMap)
|
|
|
|
defragmentWith :: ([DiskItem] -> [DiskItem]) -> [DiskItem] -> [DiskItem]
|
|
defragmentWith f diskMap
|
|
| 2 == length (groupBy (\a b -> isFile a && isFile b || isFree a && isFree b) diskMap) = diskMap
|
|
| otherwise = defragment (f diskMap)
|
|
|
|
defragment :: [DiskItem] -> [DiskItem]
|
|
defragment = defragmentWith defragmentStep
|
|
|
|
checksum :: [DiskItem] -> Int
|
|
checksum = checksum' 0
|
|
where
|
|
checksum' :: Int -> [DiskItem] -> Int
|
|
checksum' _ [] = 0
|
|
checksum' c ((Free len) : rest) = checksum' (c + len) rest
|
|
checksum' c ((File fId len) : rest) = fId * sum (take len [c..]) + checksum' (c + len) rest
|
|
|
|
part1 :: String -> IO ()
|
|
part1 s =
|
|
let diskMap = (expandDiskMap . readDiskMap) s
|
|
in print $ checksum $ defragment diskMap
|
|
|
|
-- | Part 2
|
|
freeCandidates :: [DiskItem] -> [(DiskItem, [DiskItem])]
|
|
freeCandidates [] = []
|
|
freeCandidates ((File f len) : ds) = (File f len, []) : freeCandidates ds
|
|
freeCandidates ((Free len) : ds) = (Free len, filter ((&&) <$> ((len >=) . fileLen) <*> isFile) ds) : freeCandidates ds
|
|
|
|
replaceItem :: Int -> [DiskItem] -> [DiskItem]
|
|
replaceItem _ [] = []
|
|
replaceItem fId (Free n : ds) = Free n : replaceItem fId ds
|
|
replaceItem fId (File f len : ds)
|
|
| fId == f = Free len : replaceItem fId ds
|
|
| otherwise = File f len : replaceItem fId ds
|
|
|
|
defragmentStep2 :: [DiskItem] -> [DiskItem]
|
|
defragmentStep2 diskMap =
|
|
let candidates = freeCandidates diskMap
|
|
replacePos = findIndex (not . null . snd) candidates
|
|
in case replacePos of
|
|
Nothing -> diskMap
|
|
Just pos ->
|
|
let replacement = (last . snd) (candidates !! pos)
|
|
freeLen = (fileLen . fst) (candidates !! pos)
|
|
in take pos diskMap
|
|
++ [replacement]
|
|
++ ([Free (freeLen - fileLen replacement) | (freeLen - fileLen replacement) /= 0])
|
|
++ take (length diskMap - pos) (drop (pos + 1) (replaceItem (fileId replacement) diskMap))
|
|
|
|
defragment2 :: [DiskItem] -> [DiskItem]
|
|
defragment2 diskMap =
|
|
let stepped = defragmentStep2 diskMap
|
|
in if stepped == diskMap then diskMap else defragment2 stepped
|
|
|
|
part2 :: String -> IO ()
|
|
part2 s =
|
|
let diskMap = readDiskMap s
|
|
in print $ checksum $ defragment2 diskMap
|
|
|
|
-- | Main
|
|
main :: IO ()
|
|
main = readFile "resources/day9.txt" >>= part2
|