Kattis Infinite 2D Array

infinite2darray.hs

{-# LANGUAGE BangPatterns, TypeApplications #-}
import Data.Array.Unboxed
import Data.List

-- Solution to Infinite 2D Array
-- https://open.kattis.com/problems/infinite2darray

main :: IO ()
main = do
    [x, y] <- map read . words <$> getLine
    print $ solve x y

solve :: Int -> Int -> Int
solve 0 y | y > 0 = solve y 0
solve x y = foldl' plusmod (fib (x + 2 * y)) (map contrib [1..y]) where
    contrib i = (fib i `minusmod` fib (2 * i)) `mulmod` binom (y - i + x - 1) (y - i)

mx :: Int
mx = 3000000

fib :: Int -> Int
fib = (fibs!) where
    fibs = listArray @UArray (0, mx) $ unfoldr (\(!a, !b) -> Just (a, (b, plusmod a b))) (0, 1)

binom :: Int -> Int -> Int
binom = go where
    go n k | k < 0 || n < k = 0
    go n k = fac!n `mulmod` ifac!k `mulmod` ifac!(n - k)
    fac = listArray @UArray (0, mx) $ unfoldr (\(!i, !x) -> Just (x, (i + 1, mulmod i x))) (1, 1)
    ifac = array @UArray (0, mx) $ unfoldr (\(!i, !x) -> if i < 0 then Nothing else Just ((i, x), (i - 1, mulmod x i))) (mx, inv (fac!mx))

mm :: Int
mm = 1000000007

infixl 6 `plusmod`
infixl 6 `minusmod`
infixl 7 `mulmod`
plusmod, minusmod, mulmod :: Int -> Int -> Int
plusmod a b = let c = a + b in if c >= mm then c - mm else c
minusmod a b = let c = a - b in if c < 0 then c + mm else c
mulmod a b = mod (a * b) mm

inv :: Int -> Int
inv x = go x (mm - 2) where
    go _ 0 = 1
    go x y | even y    = go (mulmod x x) (div y 2)
           | otherwise = mulmod x (go x (y - 1))

infinite2darray_formula.md

$$ F_{x,y} = f_{x+2y} + \sum_{i=1}^y (f_i - f_{2i}) \binom{y-i+x-1}{y-i} $$