{-# LANGUAGE BangPatterns #-} {-# LANGUAGE NumericUnderscores #-} import Control.Monad import Control.Monad.Cont import Data.IORef import qualified Data.Vector.Fusion.Stream.Monadic as VFSM zeta :: Int -> Int -> IO Double zeta a x = do res <- newIORef (a - 1) withBreak $ \break -> do rep (x - 1) $ \_ -> do liftIO $ modifyIORef res (* a) ret <- liftIO $ readIORef res when (ret > 1_000_000_000_000) $ break () d <- readIORef res return (1.0 / fromIntegral d) stream :: Monad m => Int -> Int -> VFSM.Stream m Int stream !l !r = VFSM.Stream step l where step x | x < r = return $ VFSM.Yield x (x + 1) | otherwise = return $ VFSM.Done {-# INLINE [0] step #-} {-# INLINE [1] stream #-} rep :: Monad m => Int -> (Int -> m ()) -> m () rep n = flip VFSM.mapM_ (stream 0 n) {-# INLINE rep #-} rep2 :: Monad m => (Int -> m ()) -> m () rep2 = flip VFSM.mapM_ (stream 2 1_000_000) withBreak :: ((r -> ContT r IO b) -> ContT r IO r) -> IO r withBreak = flip runContT pure . callCC {-# INLINE withBreak #-} main :: IO () main = do x <- readLn :: IO Int if x == 2 then putStrLn "1000000" else do res <- newIORef (0.0 :: Double) rep2 $ \a -> do d <- zeta a x modifyIORef' res (+ d) print . floor . (* 1000000) =<< readIORef res