Competitive programming using Haskell

I am currently trying to refresh my Haskell knowledge by solving some Hackerrank problems.

For example:

https://www.hackerrank.com/challenges/maximum-palindromes/problem

I've already implemented an imperative solution in C++ which got accepted for all test cases. Now I am trying to come up with a pure functional solution in (reasonably idiomatic) Haskell.

My current code is

module Main where

import           Control.Monad
import qualified Data.ByteString.Char8 as C
import           Data.Bits
import           Data.List
import qualified Data.Map.Strict       as Map
import qualified Data.IntMap.Strict    as IntMap
import           Debug.Trace

-- precompute factorials
compFactorials :: Int -> Int -> IntMap.IntMap Int
compFactorials n m = go 0 1 IntMap.empty
  where
    go a acc map
      | a < 0     = map
      | a < n     = go a' acc' map'
      | otherwise = map'
      where
        map' = IntMap.insert a acc map
        a'   = a + 1
        acc' = (acc * a') `mod` m

-- precompute invs
compInvs :: Int -> Int -> IntMap.IntMap Int -> IntMap.IntMap Int
compInvs n m facts = go 0 IntMap.empty
  where
    go a map
      | a < 0     = map
      | a < n     = go a' map'
      | otherwise = map'
      where
        map' = IntMap.insert a v map
        a' = a + 1
        v = (modExp b (m-2) m) `mod` m
        b = (IntMap.!) facts a


modExp :: Int -> Int -> Int -> Int
modExp b e m = go b e 1
  where
    go b e r
      | (.&.) e 1 == 1 = go b' e' r'
      | e > 0 = go b' e' r
      | otherwise = r
        where
          r' = (r * b) `mod` m
          b' = (b * b) `mod` m
          e' = shift e (-1)

-- precompute frequency table
initFreqMap :: C.ByteString -> Map.Map Char (IntMap.IntMap Int)
initFreqMap inp = go 1 map1 map2 inp
  where
    map1 = Map.fromList $ zip ['a'..'z'] $ repeat 0
    map2 = Map.fromList $ zip ['a'..'z'] $ repeat IntMap.empty

    go idx m1 m2 inp
      | C.null inp = m2
      | otherwise  = go (idx+1) m1' m2' $ C.tail inp
      where
        m1' = Map.update (\v -> Just $ v+1) (C.head inp) m1
        m2' = foldl' (\m w -> Map.update (\v -> liftM (\c -> IntMap.insert idx c v) $ Map.lookup w m1') w m)
              m2 ['a'..'z']


query :: Int -> Int -> Int -> Map.Map Char (IntMap.IntMap Int)
         -> IntMap.IntMap Int -> IntMap.IntMap Int -> Int
query l r m freqMap facts invs
  | x > 1     = (x * y) `mod` m
  | otherwise = y
  where
    calcCnt cs = cr - cl
      where
         cl = IntMap.findWithDefault 0 (l-1) cs
         cr = IntMap.findWithDefault 0 r cs

    f1 acc cs
      | even cnt = acc
      | otherwise = acc + 1
      where
        cnt = calcCnt cs

    f2 (acc1,acc2) cs
      | cnt < 2   = (acc1 ,acc2)
      | otherwise = (acc1',acc2')
      where
        cnt = calcCnt cs

        n = cnt `div` 2

        acc1' = acc1 + n
        r = choose acc1' n
        acc2' = (acc2 * r) `mod` m


    -- calc binomial coefficient using Fermat's little theorem
    choose n k
      | n < k = 0
      | otherwise = (f1 * t) `mod` m
      where
        f1 = (IntMap.!) facts n
        i1 = (IntMap.!) invs k
        i2 = (IntMap.!) invs (n-k)

        t = (i1 * i2) `mod` m


    x = Map.foldl' f1 0 freqMap
    y = snd $ Map.foldl' f2 (0,1) freqMap


main :: IO()
main = do
    inp <- C.getLine
    q   <- readLn :: IO Int

    let modulo  = 1000000007
    let facts   = compFactorials (C.length inp) modulo
    let invs    = compInvs (C.length inp) modulo facts
    let freqMap = initFreqMap inp

    forM_ [1..q] $ \_ -> do

      line <- getLine

      let [s1, s2] = words line
      let l = (read s1) :: Int
      let r = (read s2) :: Int

      let result = query l r modulo freqMap facts invs

      putStrLn $ show result

It passes all small and medium test cases but I am getting timeout with large test cases. The key to solve this problem is to precompute some stuff once at the beginning and use them to answer the individual queries efficiently.

Now, my main problem where I need help is:

The initital profiling shows that the lookup operation of the IntMap seems to be the main bottleneck. Is there better alternative to IntMap for memoization? Or should I look at Vector or Array, which I believe will lead to more "ugly" code. Even in current state, the code doesn't look nice (by functional standards) and as verbose as my C++ solution. Any tips to make it more idiomatic? Other than IntMap usage for memoization, do you spot any other obvious problems which can lead to performance problems?

And is there any good sources, where I can learn how to use Haskell more effectively for competitive programming?

A sample large testcase, where the current code gets timeout:

input.txt output.txt

For comparison my C++ solution:

#include <vector>
#include <iostream>

#define MOD 1000000007L

long mod_exp(long b, long e) {
    long r = 1;

    while (e > 0) {
        if ((e & 1) == 1) {
            r = (r * b) % MOD;
        }

        b = (b * b) % MOD;
        e >>= 1;
    }

    return r;
}

long n_choose_k(int n, int k, const std::vector<long> &fact_map, const std::vector<long> &inv_map) {
    if (n < k) {
        return 0;
    }

    long l1 = fact_map[n];
    long l2 = (inv_map[k] * inv_map[n-k]) % MOD;

    return (l1 * l2) % MOD;
}

int main() {
    std::string s;
    int q;

    std::cin >> s >> q;

    std::vector<std::vector<long>> freq_map;
    std::vector<long> fact_map(s.size()+1);
    std::vector<long> inv_map(s.size()+1);

    for (int i = 0; i < 26; i++) {
        freq_map.emplace_back(std::vector<long>(s.size(), 0));
    }

    std::vector<long> acc_map(26, 0);
    for (int i = 0; i < s.size(); i++) {
        acc_map[s[i]-'a']++;

        for (int j = 0; j < 26; j++) {
            freq_map[j][i] = acc_map[j];
        }
    }

    fact_map[0] = 1;
    inv_map[0] = 1;
    for (int i = 1; i <= s.size(); i++) {
        fact_map[i] = (i * fact_map[i-1]) % MOD;
        inv_map[i] = mod_exp(fact_map[i], MOD-2) % MOD;
    }

    while (q--) {
        int l, r;

        std::cin >> l >> r;
        std::vector<long> x(26, 0);

        long t = 0;
        long acc = 0;
        long result = 1;

        for (int i = 0; i < 26; i++) {
            auto cnt = freq_map[i][r-1] - (l > 1 ? freq_map[i][l-2] : 0);

            if (cnt % 2 != 0) {
                t++;
            }

            long n = cnt / 2;

            if (n > 0) {
                acc += n;
                result *= n_choose_k(acc, n, fact_map, inv_map);
                result = result % MOD;
            }
        }

        if (t > 0) {
            result *= t;
            result = result % MOD;
        }

        std::cout << result << std::endl;
    }
}

UPDATE:

DanielWagner's answer has confirmed my suspicion that the main problem in my code was the usage of IntMap for memoization. Replacing IntMap with Array made my code perform similar to DanielWagner's solution.

module Main where

import           Control.Monad
import           Data.Array            (Array)
import qualified Data.Array            as A
import qualified Data.ByteString.Char8 as C
import           Data.Bits
import           Data.List
import           Debug.Trace


-- precompute factorials
compFactorials :: Int -> Int -> Array Int Int
compFactorials n m = A.listArray (0,n) $ scanl' f 1 [1..n]
  where
    f acc a = (acc * a) `mod` m

-- precompute invs
compInvs :: Int -> Int -> Array Int Int -> Array Int Int
compInvs n m facts = A.listArray (0,n) $ map f [0..n]
  where
    f a = (modExp ((A.!) facts a) (m-2) m) `mod` m

modExp :: Int -> Int -> Int -> Int
modExp b e m = go b e 1
  where
    go b e r
      | (.&.) e 1 == 1 = go b' e' r'
      | e > 0 = go b' e' r
      | otherwise = r
        where
          r' = (r * b) `mod` m
          b' = (b * b) `mod` m
          e' = shift e (-1)

-- precompute frequency table
initFreqMap :: C.ByteString -> Map.Map Char (Array Int Int)
initFreqMap inp = Map.fromList $ map f ['a'..'z']
  where
    n = C.length inp
    f c = (c, A.listArray (0,n) $ scanl' g 0 [0..n-1])
      where
        g x j
          | C.index inp j == c = x+1
          | otherwise = x

query :: Int -> Int -> Int -> Map.Map Char (Array Int Int)
         -> Array Int Int -> Array Int Int -> Int
query l r m freqMap facts invs
  | x > 1     = (x * y) `mod` m
  | otherwise = y
  where
    calcCnt freqMap = cr - cl
      where
         cl = (A.!) freqMap (l-1)
         cr = (A.!) freqMap r

    f1 acc cs
      | even cnt = acc
      | otherwise = acc + 1
      where
        cnt = calcCnt cs

    f2 (acc1,acc2) cs
      | cnt < 2   = (acc1 ,acc2)
      | otherwise = (acc1',acc2')
      where
        cnt = calcCnt cs

        n = cnt `div` 2

        acc1' = acc1 + n
        r = choose acc1' n
        acc2' = (acc2 * r) `mod` m


    -- calc binomial coefficient using Fermat's little theorem
    choose n k
      | n < k = 0
      | otherwise = (f1 * t) `mod` m
      where
        f1 = (A.!) facts n
        i1 = (A.!) invs k
        i2 = (A.!) invs (n-k)

        t = (i1 * i2) `mod` m


    x = Map.foldl' f1 0 freqMap
    y = snd $ Map.foldl' f2 (0,1) freqMap


main :: IO()
main = do
    inp <- C.getLine
    q   <- readLn :: IO Int

    let modulo  = 1000000007
    let facts   = compFactorials (C.length inp) modulo
    let invs    = compInvs (C.length inp) modulo facts
    let freqMap = initFreqMap inp

    replicateM_ q $ do

      line <- getLine

      let [s1, s2] = words line
      let l = (read s1) :: Int
      let r = (read s2) :: Int

      let result = query l r modulo freqMap facts invs

      putStrLn $ show result