NextPalindrome.hs

import Control.Applicative
import qualified Data.ByteString.Char8 as B
import Control.Monad.ST
import Data.STRef
import Control.Monad
import Data.Array.ST

addOne :: B.ByteString -> B.ByteString
addOne xs =
        let fix ('a':rest) = '1':'0':rest
            fix x = x
            l  = B.length xs
        in  B.pack . fix . runST $
          do arr <- newListArray (1, l) (B.unpack xs) :: ST s (STArray s Int Char)
             -- arr <- newArray_ (1, l)  :: ST s (STArray s Int Char)
             -- forM_ [1..l] (\idx -> writeArray arr idx (xs `B.index` (idx - 1)))
             let loop 0   = writeArray arr 1 'a'
                 loop idx = do
                    value <- readArray arr idx
                    case value of
                     '9' -> writeArray arr idx '0' >> loop (idx - 1)
                     x   -> writeArray arr idx (succ x)
             loop l
             getElems arr

biggerOrEquals :: B.ByteString -> B.ByteString -> Bool
biggerOrEquals x y =
         let (lx, ly) = (B.length x, B.length y)
         in case compare lx ly of
            GT -> True
            LT -> False
            EQ -> case dropWhile (==EQ) $ B.zipWith compare x y of
                  [] -> True
                  (GT:_) -> True
                  _  -> False

-- addOne :: B.ByteString -> B.ByteString
-- addOne = B.reverse . inc . B.reverse
--        where inc x | B.length x == 0 = B.singleton '1'
--                    | B.head x == '9' =  '0' `B.cons` (inc $ B.tail x)
--                    | otherwise = (succ $ B.head x) `B.cons` (B.tail x)

-- addOne :: B.ByteString -> B.ByteString
-- addOne = B.reverse . B.pack . inc . B.unpack . B.reverse
--        where inc x | length x == 0 = ['1']
--                    | head x == '9' =  '0' : (inc $ tail x)
--                    | otherwise = (succ $ head x) : (tail x)

solve :: B.ByteString -> B.ByteString
solve x | B.length x == 0 = error "empty list"
        | B.length x == 1 = x
        | otherwise =
      let half = B.length x `div` 2
          (l, rs) = B.splitAt half x
          rl     = B.reverse l
      in if even $ B.length x then
            if rl `biggerOrEquals` rs
            then l `B.append` rl
            else solve $ (addOne l) `B.append` (B.replicate half '0')
         else let (fr,r) = (B.head rs, B.tail rs)
              in if rl `biggerOrEquals` r
                 then l `B.append` (fr `B.cons` rl)
                 else solve $ (addOne (l `B.snoc` fr)) `B.append` (B.replicate half '0')

output :: B.ByteString -> IO ()
output = B.putStrLn

main :: IO ()
main = do
     n:ls <- B.lines <$> B.getContents
--     B.interact (B.unlines . map (solve . addOne) . tail . (B.lines))
     mapM_ (output . solve . addOne) $ take (read $ B.unpack n) ls
     return ()