#579013 (Haskell) No.1240 Or Sum of Xor Pair

• 問題: No.1240 Or Sum of Xor Pair
• ユーザー: こまる
• 提出日時: 2020-11-13 14:27:02
• 言語: Haskell (9.10.1)
• 結果: CE (最新)
  AC
  (最初)
• 実行時間: -
• コード長: 5,675 bytes
• コンパイル時間: 565 ms
• コンパイル使用メモリ: 202,880 KB
• 最終ジャッジ日時: 2024-11-14 23:54:19
• 合計ジャッジ時間: 2,027 ms
• ジャッジサーバーID （参考情報）: judge2 / judge4

コンパイルメッセージ

Loaded package environment from /home/judge/.ghc/x86_64-linux-9.8.2/environments/default
[1 of 2] Compiling Main             ( Main.hs, Main.o )

Main.hs:24:25: error: [GHC-87543]
    Ambiguous occurrence ‘.<<.’.
    It could refer to
       either ‘Data.Bits..<<.’,
              imported from ‘Data.Bits’ at Main.hs:9:1-26,
           or ‘Main..<<.’, defined at Main.hs:154:1.
   |
24 |   f <- VUM.replicate (1 .<<. m) (0 :: Int)
   |                         ^^^^

Main.hs:25:25: error: [GHC-87543]
    Ambiguous occurrence ‘.<<.’.
    It could refer to
       either ‘Data.Bits..<<.’,
              imported from ‘Data.Bits’ at Main.hs:9:1-26,
           or ‘Main..<<.’, defined at Main.hs:154:1.
   |
25 |   g <- VUM.replicate (1 .<<. m) (0 :: Int)
   |                         ^^^^

Main.hs:35:28: error: [GHC-87543]
    Ambiguous occurrence ‘.<<.’.
    It could refer to
       either ‘Data.Bits..<<.’,
              imported from ‘Data.Bits’ at Main.hs:9:1-26,
           or ‘Main..<<.’, defined at Main.hs:154:1.
   |
35 |     ff <- VUM.replicate (1 .<<. m) (0 :: Int)
   |                            ^^^^

Main.hs:36:50: error: [GHC-87543]
    Ambiguous occurrence ‘.>>.’.
    It could refer to
       either ‘Data.Bits..>>.’,
              imported from ‘Data.Bits’ at Main.hs:9:1-26,
           or ‘Main..>>.’, defined at Main.hs:156:1.
   |
36 |     VU.forM_ an $ \a -> VUM.unsafeModify ff (+(a .>>. d .&. 1)) a
   |                                                  ^^^^

Main.hs:40:40: error: [GHC-87543]
    Ambiguous occurrence ‘.<<.’.
    It could refer to
       either ‘Data.Bits..<<.’,
              imported from ‘Data.Bits’ at Main.hs:9:1-26,
           or ‘Main..<<.’, defined at Main.hs:154:1.
   |
40 |     modifyIORef' ans (subtract (ff2Sum .<<. d))
   |                                        ^^^^

Main.hs:142:46: error: [GHC-87543]
    Ambiguous occurrence ‘.<<.’.
    It could refer to
 

ソースコード

{-# LANGUAGE BangPatterns         #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE MagicHash            #-}
{-# LANGUAGE OverloadedStrings    #-}
{-# LANGUAGE UndecidableInstances #-}

import           Control.Monad
import           Control.Monad.State
import           Data.Bits
import           Data.Char
import           Data.IORef
import           GHC.Exts
import qualified Data.ByteString.Char8             as BSC8
import qualified Data.Vector.Fusion.Stream.Monadic as VFSM
import qualified Data.Vector.Unboxed               as VU
import qualified Data.Vector.Unboxed.Mutable       as VUM

main :: IO ()
main = do
  [n, x] <- map (read :: String -> Int) . words <$> getLine
  an <- seqInput n
  let
    m = bitLength $ VU.maximum an
  f <- VUM.replicate (1 .<<. m) (0 :: Int)
  g <- VUM.replicate (1 .<<. m) (0 :: Int)
  VU.forM_ an $ \a -> do
    VUM.unsafeModify f (+a) a 
    VUM.unsafeModify g (+1) a
  f' <- VU.unsafeFreeze f
  g' <- VU.unsafeFreeze g
  let fg = convolute f' g' XOR
  let fgSum = (*2) . VU.sum . VU.take x $ fg
  ans <- newIORef fgSum
  rep m $ \d -> do
    ff <- VUM.replicate (1 .<<. m) (0 :: Int)
    VU.forM_ an $ \a -> VUM.unsafeModify ff (+(a .>>. d .&. 1)) a
    ff' <- VU.unsafeFreeze ff
    let ff2 = convolute ff' ff' XOR
    let ff2Sum = VU.sum . VU.take x $ ff2
    modifyIORef' ans (subtract (ff2Sum .<<. d))
  modifyIORef' ans ((`div` 2) . subtract (VU.sum an))
  print =<< readIORef ans


data Kinds = AND | OR | XOR
  deriving Eq

walshSpectre :: VU.Vector Int ->　VU.Vector Int
walshSpectre vec = fwt vec False False True

convolute :: VU.Vector Int -> VU.Vector Int -> Kinds -> VU.Vector Int
convolute xs ys k
  | k == AND = convoluteAND xs ys
  | k == OR  = convoluteOR xs ys
  | otherwise = convoluteXOR xs ys

convoluteAND :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int
convoluteAND xs ys =
  let
    xs' = fwt xs False True False
    ys' = fwt ys False True False
    zs' = VU.zipWith (*) xs' ys'
  in
    VU.take m $ fwt zs' True True False
  where
    !m = min (VU.length xs) (VU.length ys)

convoluteOR :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int
convoluteOR xs ys =
  let
    xs' = fwt xs False False False
    ys' = fwt ys False False False
    zs' = VU.zipWith (*) xs' ys'
  in
    fwt zs' True False False

convoluteXOR :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int
convoluteXOR xs ys =
  let
    xs' = fwt xs False False True
    ys' = fwt ys False False True
    zs' = VU.zipWith (*) xs' ys'
  in
    fwt zs' True False True

fwt :: VU.Vector Int -> Bool -> Bool -> Bool -> VU.Vector Int
fwt f' inv isAND isXOR = VU.create $ do
  g <- VU.unsafeThaw f
  rep' n $ \i -> rep n $ \j -> when (j .&. i == 0) $ do
    if isXOR
      then do
        itemX <- VUM.unsafeRead g j
        itemY <- VUM.unsafeRead g (j .|. i)
        if inv
          then do
            VUM.unsafeWrite g j ((itemX + itemY) `div` 2)
            VUM.unsafeWrite g (j .|. i) ((itemX - itemY) `div` 2)
          else do
            VUM.unsafeWrite g j (itemX + itemY)
            VUM.unsafeWrite g (j .|. i) (itemX - itemY)
      else do
        if isAND
          then do
            item <- VUM.unsafeRead g (j .|. i)
            if inv
              then VUM.unsafeModify g (subtract item) j
              else VUM.unsafeModify g (+ item) j
          else do
            item <- VUM.unsafeRead g j
            if inv
              then VUM.unsafeModify g (subtract item) (j .|. i)
              else VUM.unsafeModify g (+ item) (j .|. i)
  return g
  where
    !f = growVU f'
    !n = VU.length f

growVU :: VU.Vector Int -> VU.Vector Int
growVU v
  | VU.null v = VU.singleton 0
  | VU.length v == 1 = v
  | otherwise = v VU.++ VU.replicate (ceilPow2 n - n) 0
  where !n = VU.length v

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 #-}

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' 1 n)
{-# INLINE rep' #-}

infixl 8 .<<., .>>., .>>>.
infixl 6 .^.
(.<<.), (.>>.) :: Bits b => b -> Int -> b
(.<<.) = unsafeShiftL
{-# INLINE (.<<.) #-}
(.>>.) = unsafeShiftR
{-# INLINE (.>>.) #-}
(.>>>.) :: Int -> Int -> Int
(.>>>.) (I# x#) (I# i#) = I# (uncheckedIShiftRL# x# i#)
{-# INLINE (.>>>.) #-}
(.^.) :: Bits b => b -> b -> b
(.^.)  = xor
{-# INLINE (.^.)  #-}
clz :: FiniteBits fb => fb -> Int
clz = countLeadingZeros
{-# INLINE clz #-}
ctz :: FiniteBits fb => fb -> Int
ctz = countTrailingZeros
{-# INLINE ctz #-}
ceilPow2 :: Int -> Int
ceilPow2 n
  | n > 1     = (-1) .>>>. clz (n - 1) + 1
  | otherwise = 1
{-# INLINE ceilPow2 #-}
floorPow2 :: Int -> Int
floorPow2 n
  | n >= 1    = 1 .<<. (63 - clz n)
  | otherwise = 0
{-# INLINE floorPow2 #-}
bitLength :: FiniteBits fb => fb -> Int
bitLength n = 64 - countLeadingZeros n
{-# INLINE bitLength #-}


type CParser a = StateT BSC8.ByteString Maybe a
runCParser :: CParser a -> BSC8.ByteString -> Maybe (a, BSC8.ByteString)
runCParser = runStateT
{-# INLINE runCParser #-}
int :: CParser Int
int = coerce $ BSC8.readInt . BSC8.dropWhile isSpace
{-# INLINE int #-}
seqInput :: Int -> IO (VU.Vector Int)
seqInput n = VU.unfoldrN n (runCParser int) <$> BSC8.getLine
{-# INLINE seqInput #-}