Submission #229428

---

Source Code Expand

{-# OPTIONS_GHC -O2 -funbox-strict-fields #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE ExplicitForAll #-}

import Control.Applicative
import Data.Bits
import qualified Data.ByteString.Char8 as BS
import Data.Function
import Data.List
import Data.Maybe
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM
import Data.Word

main :: IO ()
main = do
  [n, x] <- getInts
  as <- getIntsN n
  print $ solve as x

solve :: U.Vector Int -> BitSet -> Int
solve vs s = sum $ map (vs U.!) $ bitSetToList s

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

getInts :: IO [Int]
getInts = readInts <$> BS.getLine

readInts :: BS.ByteString -> [Int]
readInts s0 = unfoldr step s0
  where
    step (BS.dropWhile (==' ') -> s)
      | s == "" = Nothing
      | Just (v, r) <- BS.readInt s = Just (v, r)
      | otherwise = error $ "not an integer: " ++ show s

getIntsN :: Int -> IO (U.Vector Int)
getIntsN n = readIntsN n <$> BS.getLine

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

----------------------------------------------------------------------------
-- BitSet

type BitSet = Int

bitSetToList :: BitSet -> [Int]
bitSetToList = unfoldr step . fromIntegral
  where
    step 0 = Nothing
    step s = Just (b, s')
      where
        !b = bitScanForward s
        !s' = s .&. complement (1 `unsafeShiftL` b)

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

bitScanForward :: Word -> Int
bitScanForward w = bsfTable U.! bsfHash (w .&. (-w))

bsfTable :: U.Vector Int
bsfTable = U.create $ do
  mv <- UM.new 64
  U.forM_ (U.generate 64 id) $ \i ->
    UM.write mv (bsfHash $ 1 `shiftL` i) i
  return mv

bsfHash :: Word -> Int
bsfHash w = fromIntegral $ (w * 0x03f79d71b4cb0a89) `shiftR` 58

Submission Info

Judge Result