|
| 1 | +-- https://github.com/minoki/my-atcoder-solutions |
| 2 | +{-# LANGUAGE BangPatterns #-} |
| 3 | +import Data.Char (isSpace) |
| 4 | +import Data.Int (Int64) |
| 5 | +import Data.List (unfoldr) |
| 6 | +import qualified Data.Vector.Unboxed as U |
| 7 | +import qualified Data.ByteString.Char8 as BS |
| 8 | +import qualified Data.IntSet as IntSet |
| 9 | +import Control.Exception (assert) |
| 10 | +import qualified Test.QuickCheck as QC |
| 11 | +import GHC.Stack (HasCallStack) |
| 12 | + |
| 13 | +(!) :: (HasCallStack, U.Unbox a) => U.Vector a -> Int -> a |
| 14 | +(!) = (U.!) |
| 15 | +{- |
| 16 | +(!) vec i | i < 0 = error $ "negative index: " ++ show i |
| 17 | + | i >= U.length vec = error $ "out of bounds " ++ show (i,U.length vec) |
| 18 | + | otherwise = vec U.! i |
| 19 | +-} |
| 20 | + |
| 21 | +cycles :: Int -> U.Vector Int -> [[Int]] |
| 22 | +cycles !n perm = loop (IntSet.fromDistinctAscList [0..n-1]) |
| 23 | + where |
| 24 | + loop s | IntSet.null s = [] |
| 25 | + | otherwise = let (m,s') = IntSet.deleteFindMin s |
| 26 | + (cyc,s'') = oneCycle m m [m] s' |
| 27 | + in cyc : loop s'' |
| 28 | + oneCycle m0 m xs s = let m' = perm ! m |
| 29 | + in if m' == m0 then |
| 30 | + (xs, s) |
| 31 | + else |
| 32 | + oneCycle m0 m' (m':xs) (IntSet.delete m' s) |
| 33 | + |
| 34 | +solve :: Int -> Int -> U.Vector Int -> U.Vector Int64 -> Int64 |
| 35 | +solve !n !k perm c = maximum $ map solveOneCycle $ cycles n perm |
| 36 | + where |
| 37 | + solveOneCycle cyc = |
| 38 | + let scores = U.fromList $ map (c !) cyc |
| 39 | + cycle_len = U.length scores |
| 40 | + t = U.sum scores |
| 41 | + scores' = U.init $ scores <> scores |
| 42 | + ss = U.scanl' (+) 0 scores' |
| 43 | + !_ = assert (U.length ss == 2 * cycle_len) |
| 44 | + in if t <= 0 || k <= cycle_len then |
| 45 | + maximum [ ss ! j - ss ! i |
| 46 | + | i <- [0 .. cycle_len - 1] |
| 47 | + , j <- [i + 1 .. min (2 * cycle_len - 1) (i + k)] |
| 48 | + ] |
| 49 | + else |
| 50 | + let (q,r) = k `quotRem` cycle_len |
| 51 | + -- q >= 1 |
| 52 | + in t * fromIntegral q + if r == 0 then |
| 53 | + max 0 $ maximum [ ss ! j - ss ! i |
| 54 | + | i <- [0 .. cycle_len - 1] |
| 55 | + , j <- [i + 1 .. min (2 * cycle_len - 1) (i + cycle_len)] |
| 56 | + ] - t |
| 57 | + else |
| 58 | + maximum [ ss ! j - ss ! i |
| 59 | + | i <- [0 .. cycle_len - 1] |
| 60 | + , j <- [i + 1 .. min (2 * cycle_len - 1) (i + r)] |
| 61 | + ] |
| 62 | + |
| 63 | +main = do |
| 64 | + [n,k] <- unfoldr (BS.readInt . BS.dropWhile isSpace) <$> BS.getLine |
| 65 | + perm <- U.map (subtract 1) <$> U.unfoldrN n (BS.readInt . BS.dropWhile isSpace) <$> BS.getLine |
| 66 | + c <- U.map fromIntegral <$> U.unfoldrN n (BS.readInt . BS.dropWhile isSpace) <$> BS.getLine |
| 67 | + print $ solve n k perm c |
| 68 | + |
| 69 | +prop :: QC.Property |
| 70 | +prop = let gen = do n <- QC.choose (2, 100) |
| 71 | + -- n <- QC.choose (2, 5000) |
| 72 | + k <- QC.choose (1, 10^9) |
| 73 | + perm <- QC.shuffle [0..n-1] |
| 74 | + c <- QC.vectorOf n (QC.choose (-10^9, 10^9)) |
| 75 | + return (n, k, U.fromList perm, U.fromList c) |
| 76 | + in QC.forAll gen (\(n,k,perm,c) -> solve n k perm c `seq` ()) |
0 commit comments