zigzag.hs

import Data.List

-- 第一行照写，第二行开头加空格反写；如果不够一行，前面加空格填充
arrange :: Int -> String -> [String]
arrange n str
  | n == 1 = [str]
  | length str <= n = [str]
  | length str <= (2*n-1) = [take n str,replicate (2*n-1 - length str) ' ' ++ reverse (drop n str) ++ " "]
  | otherwise = arrange n (take (2*n-1) str) ++ arrange n (drop (2*n-1) str)

-- 从第二组开始，填充一些空格，使得arrange得以实现
fill1 n str
  | n == 1 = str
  | length str <= 2*n-1 = str
  | otherwise = (take (2*n-1) str) ++ fill2 n (drop (2*n-1) str)
  where fill2 n st
          | length st <= 2*n-2 = " "++ st
          | otherwise = " " ++ (take (2*n-2) st) ++ fill2 n (drop (2*n-2) st)

zigzag n str = filter (/= ' ') $ concat $ transpose $ arrange n (fill1 n str)

tests = [zigzag 3 "PAYPALISHIRING" == "PAHNAPLSIIGYIR", zigzag 3 "ABCDEF" == "AEBDFC", zigzag 4 "ABCDEF" == "ABFCED", zigzag 4 "ABCDE" == "ABCED", zigzag 1 "AB" == "AB", zigzag 2 "ABCD" == "ACBD", zigzag 2 "ABC" == "ACB"]

main = print $ all (== True) tests

http://oj.leetcode.com/problems/zigzag-conversion/