みんなのプロコン 2018 C問題 - 駆引取引

yahoo_procon2018_c.hs

module Main where

-- 高橋君と青木君が取引をします。 はじめ、高橋君の所持金は 0 円です。
-- 高橋君は 1 から N の番号がついた N 個の財宝を持っています。高橋君は青木君に財宝 i を売却すると xi 円得ることができます。
-- 青木君は 1 から N の番号がついた N 個の商品を販売しています。商品 i は価値 vi を持ち、価格 ci 円で販売されています。
-- 取引は以下の手順で行われます。
-- 1. 高橋君は財宝を売却するか、商品を購入するかを決める。前者ならば手順 2. へ、後者ならば手順 3. へ進む。
-- 2. 高橋君は持っている財宝のうち、最も番号が小さいものを青木君に売却しお金を得る。その後、青木君は商品を 1 つ選び販売を停止する。手順 1. へ戻る。
-- 3. 高橋君は販売中の 0 個以上の商品を購入し、取引を終了する。このとき、高橋君は所持金が 0 円未満になるように購入することはできない。
-- 高橋君が購入した商品の価値の総和を得点とします。
-- 高橋君は得点を最大化するように、青木君は得点を最小化するように行動するとき、得点はいくつになりますか？

import Data.List
import Debug.Trace
import GHC.Exts

newtype ProductSell = ProductSell { getSellPrice :: Int } deriving (Show, Eq)

data ProductBuy = ProductBuy { getPrice:: Int
                             , getValue:: Int
                             } deriving (Show, Eq)

-- 売れる個数を決めれる
-- 売れる個数ごとにその金額内で最大バリューを選択して、それを選択しから除く
-- 除かれた商品の中から、最大バリューとなる結果を選択
-- 売れる個数ごとに結果を出して最大の値を出す

-- | combination 関数
-- | ref: http://www.geocities.jp/m_hiroi/func/haskell06.html
--
-- >>> combinations 1 [1]
-- [[1]]
-- >>> combinations 1 [1..2]
-- [[1],[2]]
-- >>> combinations 2 [1..3]
-- [[1,2],[1,3],[2,3]]
-- >>> combinations 2 [1..4]
-- [[1,2],[1,3],[1,4],[2,3],[2,4],[3,4]]
combinations :: Int -> [a] -> [[a]]
combinations n xs = comb n (length xs) xs where
  comb 0 _ _ = [[]]
  comb r n a@(x:xs)
    | n == r    = [a]
    | otherwise = map (x:) (comb (r - 1) (n - 1) xs) ++ comb r (n - 1) xs

-- | 売った場合の金額合計を個数ごとに求める
--
-- >>> getAccumulatedSellPrice (map ProductSell [200, 1000, 1])
-- [(0,0),(1,200),(2,1200),(3,1201)]
-- >>> getAccumulatedSellPrice (map ProductSell [1, 100])
-- [(0,0),(1,1),(2,101)]
-- >>> getAccumulatedSellPrice (map ProductSell [89,36,77,25,84,49,67,21,78,94,55,50,22,40,44,17,93,11])
-- [(0,0),(1,89),(2,125),(3,202),(4,227),(5,311),(6,360),(7,427),(8,448),(9,526),(10,620),(11,675),(12,725),(13,747),(14,787),(15,831),(16,848),(17,941),(18,952)]
getAccumulatedSellPrice :: [ProductSell] -> [(Int, Int)]
getAccumulatedSellPrice products = [(length x, (sum . map getSellPrice) x) | x <- inits products]

-- | 指定した金額以下で買える商品の組み合わせを割り出す
--
-- >>> getAffordableProducts (zipWith ProductBuy [100, 100, 300] [1000, 2000, 1000000]) 1 99
-- []
-- >>> getAffordableProducts (zipWith ProductBuy [100, 100, 300] [1000, 2000, 1000000]) 1 100
-- [[ProductBuy {getPrice = 100, getValue = 1000}],[ProductBuy {getPrice = 100, getValue = 2000}]]
-- >>> getAffordableProducts (zipWith ProductBuy [100, 100, 300] [1000, 2000, 1000000]) 2 100
-- []
getAffordableProducts :: [ProductBuy] -> Int -> Int -> [[ProductBuy]]
getAffordableProducts products n price = [x | x <- combinations n products, (sum . map getPrice) x <= price]


-- | 販売停止する個数と、その際に購入されうる金額内で最大の価値になる組み合わせを探索する
--
-- >>> getMostValueWithPriceLimit (zipWith ProductBuy [100, 100, 300] [1000, 2000, 1000000]) 1 99
-- []
-- >>> getMostValueWithPriceLimit (zipWith ProductBuy [100, 100, 300] [1000, 2000, 1000000]) 1 100
-- [ProductBuy {getPrice = 100, getValue = 2000}]
getMostValueWithPriceLimit :: [ProductBuy] -> Int -> Int -> [ProductBuy]
getMostValueWithPriceLimit products n price | null candidate = []
                                            | otherwise = head (sortWith (Down . sum . map getValue) candidate)
                                               where
                                                     candidate = getAffordableProducts products n price

-- | ソルバー
--
-- >>> solve (map ProductSell [200, 1000, 1]) (zipWith ProductBuy [100, 100, 300] [1000, 2000, 1000000])
-- 1000
-- >>> solve (map ProductSell [1, 100]) (zipWith ProductBuy [1, 100] [100, 1])
-- 0
-- >>> solve (map ProductSell [100,1,1,1,1]) (zipWith ProductBuy [1,1,1,1,1] [1000000000000000,1000000000000000,1000000000000000,1000000000000000,1000000000000000])
-- 4000000000000000
-- >>> solve (map ProductSell [89,36,77,25,84,49,67,21,78,94,55,50,22,40,44,17,93,11]) (zipWith ProductBuy [55,81,89,38,19,65,68,75,66,52,75,93,94,31,73,86,25,79] [38,81,44,33,70,20,89,34,13,45,27,88,91,85,36,98,52,93])
-- 337
solve :: [ProductSell] -> [ProductBuy] -> Int
solve sellItems buyItems | null candidate = 0
                         | otherwise = sum (map getValue (head candidate))
  where
    priceLimits = getAccumulatedSellPrice sellItems
    itemsToStopToSell = [(n, price, getMostValueWithPriceLimit buyItems n price) | (n, price) <- priceLimits]
    candidate = sortWith (Down . sum . map getValue) [getMostValueWithPriceLimit (buyItems \\ stopped) (length buyItems - n) price | (n, price, stopped) <- itemsToStopToSell]

main :: IO ()
main = do
  n <- read <$> getLine :: IO Int
  xs <- (map read . words) <$> getLine :: IO [Int]
  cs <- (map read . words) <$> getLine :: IO [Int]
  vs <- (map read . words) <$> getLine :: IO [Int]
  let productSell = map ProductSell xs
  let productBuy = zipWith ProductBuy cs vs
  print (solve productSell productBuy)