Optimize calculation of string self-similarly with its suffixes

Question

I am trying to solve a Hacker Rank problem about string suffixes:

For two strings A and B, we define the similarity of the strings to be the length of the longest prefix common to both strings. For example, the similarity of strings "abc" and "abd" is 2, while the similarity of strings "aaa" and "aaab" is 3.

Calculate the sum of similarities of a string S with each of its suffixes.

I have the following code:

import System.IO
import Control.Monad
import Data.Maybe
import Data.List
import qualified Data.ByteString.Char8 as BSC

main = do
    n <- liftM getIntFromBS BSC.getLine
    replicateM_ n $ do
        s <- BSC.getLine
        putStrLn . show $ sum $ map (cntPrefix s) $ tail $ sort $ BSC.tails s
    where
        getIntFromBS = fst.fromJust.BSC.readInt
        cntPrefix str pref = length $ takeWhile (\t -> fst t == snd t) $ BSC.zip pref str

But performance is really bad. I am doing this challenges to learn Haskell and thus my skills aren't good enough to optimize this program. Any help is appreciated.

Answer 1

As far as I can tell your solution is not only slow but incorrect because of your usage of sort. Dropping that single function from your pipeline nets an exponential improvement, and it stopped including the similarity of s to itself in the answer.

Make sure also you're compiling with -O2, otherwise sum won't be optimized into a strict fold and you'll blow up the stack with thunks.

I also came up with a version targeted toward readability, but it's ~3x slower than your (corrected) ByteString based version so I'll just include it here as a curiosity.

import Control.Monad (replicateM_)
import Data.List (tails)

main :: IO ()
main = do
        n <- readLn
        replicateM_ n $ do
            s <- getLine
            print $ sum $ map (similarity s) (suffixes s)

similarity :: (Eq a) => [a] -> [a] -> Int
similarity s = length . takeWhile id . zipWith (==) s

suffixes :: [a] -> [[a]]
suffixes = tail . tails