import Control.Monad import Data.Function import Data.List import Data.Maybe type Ring = Int type Index = Int type Value = Int righttop :: Ring -> Index righttop ring = 1 + (2 * ring - 1) * (2 * ring) -- rt(r) = 1 + 2r(2r-1) = 4r^2 - 2r + 1 = n -- 4r^2 - 2r + 1 - n = 0 -- r = (2 + sqrt(4 - 4*4*(1-n))) / 8 -- = (2 + sqrt(16n - 12)) / 8 -- = (2 + 2sqrt(4n - 3)) / 8 -- = (1 + sqrt(4n - 3)) / 4 -- counted from right-top numbers (1 is also a right-top number) ringof :: Index -> Ring ringof n = floor $ (1 + sqrt (fromIntegral (4 * n - 3))) / 4 posof :: Index -> (Int, Int) posof n = let r = ringof n rt = righttop r d = 2 * r in case n of _ | n < rt + d -> (r - (n - rt), r) | n < rt + 2*d -> (-r, r - (n - (rt + d))) | n < rt + 2*d + d+1 -> (-r + n - (rt + 2*d), -r) | n < rt + 2*d + 2*(d+1) -> (r + 1, -r + n - (rt + 2*d + d+1)) | otherwise -> undefined posidx :: (Int, Int) -> Index posidx (x, y) = let ring = if x > abs y then x - 1 else (max `on` abs) x y rt = righttop ring ringlen = 8 * ring + 2 Just ringidx = findIndex (== (x, y)) (map posof [rt .. rt+ringlen-1]) in rt + ringidx manhattan1 :: Index -> Int manhattan1 = uncurry ((+) `on` abs) . posof valueof :: Index -> Value valueof n = cache !! (n - 1) where cache = map valueof' [1..] valueof' :: Index -> Value valueof' 1 = 1 valueof' n = let (x, y) = posof n nei = [(x + dx, y + dy) | dx <- [-1..1], dy <- [-1..1], dx /= 0 || dy /= 0] indices = filter (< n) (map posidx nei) in sum (map valueof indices) main :: IO () main = do input <- liftM read (readFile "3.in") :: IO Int print (manhattan1 input) -- putStrLn $ intercalate "\n" [intercalate " " $ map show [valueof (posidx (x,-y)) | x <- [-3..3]] | y <- [-3..3]] print $ fromJust $ find (> input) $ map valueof [1..]