Submission #421856

---

Source Code Expand

{-# OPTIONS_GHC -O2 -funbox-strict-fields #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE NumDecimals #-}
{-# LANGUAGE DisambiguateRecordFields #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ExplicitForAll #-}

import Control.Applicative
import Data.Bits
import qualified Data.ByteString.Char8 as BS
import qualified Data.Vector.Unboxed as U

main :: IO ()
main = do
  n <- getInt
  abbs <- U.replicateM n getInt2
  U.mapM_ print $ solve32 abbs

solve32 :: U.Vector (Int, Int) -> U.Vector Int
solve32 (U.unzip -> (as, bs)) =
  take_ (2 * U.length as) $
  U.drop 1 $
  convolve32 (U.cons 0 as) (U.cons 0 bs)
  where
    take_ k v
      | U.length v > k = U.take k v
      | otherwise = v U.++ U.replicate (k - U.length v) 0

convolve32 :: U.Vector Int -> U.Vector Int -> U.Vector Int
convolve32 as bs
  | U.null as || U.null bs = U.empty
convolve32 as bs = dec (enc as * enc bs :: Integer)
  where
    enc vec
      | len == 1 = fromIntegral $ U.head vec
      | otherwise = enc (U.take hlen vec) +
          enc (U.drop hlen vec) `shiftL` (hlen * 32)
      where
        !len = U.length vec
        hlen = div len 2
    dec n = U.generate (cdiv len 32) $ \i ->
      U.sum $ U.generate 32 $ \j ->
        if testBit n (32 * i + j)
          then 1 `shiftL` j
          else 0
      where
        len = findFirst (-1) 1e9 $ (==0) . shiftR n

----------------------------------------------------------------------------
-- IO

getInt :: IO Int
getInt = readInt <$> BS.getLine

readInt :: BS.ByteString -> Int
readInt s = case BS.readInt s of
  Just (r, "") -> r
  _ -> error $ "not an integer: " ++ show s

getInt2 :: IO (Int, Int)
getInt2 = readInt2 <$> BS.getLine

readInt2 :: BS.ByteString -> (Int, Int)
readInt2 s = (v0, v1)
  where
    !v = readIntsN 2 s
    !v0 = U.unsafeIndex v 0
    !v1 = U.unsafeIndex v 1

readIntsN :: Int -> BS.ByteString -> U.Vector Int
readIntsN n s0
  | U.length vec == n = vec
  | otherwise = error $ "readIntsN: expecting " ++ show n
    ++ " ints but got " ++ show (U.length vec)
  where
    vec = U.unfoldrN (n+1) step s0
    step (BS.dropWhile (==' ') -> s)
      | s == "" = Nothing
      | Just (v, r) <- BS.readInt s = Just (v, r)
      | otherwise = error $ "not an integer: " ++ show s

----------------------------------------------------------------------------
-- Util

infixl 7 `cdiv`
cdiv :: Int -> Int -> Int
cdiv a b = div (a + b - 1) b

----------------------------------------------------------------------------
-- BinSearch

findFirst :: Int -> Int -> (Int -> Bool) -> Int
findFirst lo hi f -- precondition: not (f lo) && f hi
  | hi - lo == 1 = hi
  | otherwise = let
    !mid = (lo + hi) `shiftR` 1
    in case f mid of
      True -> findFirst lo mid f
      False -> findFirst mid hi f

Submission Info

Judge Result