平面走査 (joi2024_yo2_b)

B - 買い物 2 (Shopping 2)

AtCoder Problems から偶然見た JOI 2023/2024 二次予選 問2が味わい深かったのでまとめておく。（他はまだ味わい中）

問題概要

問題文 公式
問題文 AtCoder

• 営業$M$日：日付1から$M$
• 商品$N$個：番号1から$N$、それぞれ、定価$P_i$(偶数)と割引き日$A_i$ ($2 \leq P_i, 1 \leq A_i \leq M$)
  割引き日には半額になる
• 客$Q$人：番号1から$Q$、それぞれ、来店日$T_k$、商品番号 $L_k$ ～ $R_k$ を一つずつ購入 ($1 \leq T_k \leq M, 1 \leq L_k \leq R_k \leq N$)

客1～$Q$ それぞれについて、購入金額を求めよ。

制約

• $1 \leq N \leq 2 \times 10^5$
• $1 \leq M \leq 2 \times 10^5$
• $1 \leq Q \leq 2 \times 10^5$
• $2 \leq P_i \leq 10^9$

（データが正の整数だとか、$A_i, T_k, L_k, R_k$ の範囲や $P_i$ が偶数であることなどは、問題の設定から必然なことで、「正常な入力の仕様」かつ「入力が正常であることの保証」と呼ぶべきではないかな。「制約」というより。）

シグネチャを決める。横着する。

compute :: [Int]   -- N,M,Q
        -> [[Int]] -- Pi, Ai
        -> [[Int]] -- Tk, Lk, Rk
        -> [Int]   -- 答えQ個
compute [n,m,q] pas tlrs = ...

「小課題」と称して、制約をさらに強くした小問により、解答を誘導しているようだ。
とても教育的で感動。素直に従ってみる。

小課題1

$N \leq 2,000, M \leq 2,000, Q \leq 2,000$

規模が小さいので、何の工夫もなく計算

• 客一人ずつについて　（全部でQ回）
• 商品一つずつについて（全部でN回）

しても間に合う。

compute [n,m,q] pas tlrs | n <= 2000, m <= 2000, q <= 2000 =
  [ sum [ if t == a then div p 2 else p
        | let l1 = pred l
        , p:a:_ <- take (r - l1) $ drop l1 pas ]
  | t:l:r:_ <- tlrs ]

小課題2

$M = 1$

全ての商品について割引日は1で、全ての客もその日に来店するので、全て半額で計算すればいい。
ただしNやQの条件は上限になっているので、合計の計算は累積和で$O(1)$にする。

import Data.Array

compute [n,1,_q] pas tlrs =
    [(pacc ! r) - (pacc ! pred l) | t:l:r:_ <- tlrs]
  where
    pacc = mkAcc n [div p 2 | p:_ <- pas]

-- mkAcc n xs :: n == length xs の累積和 添え字1始まり
mkAcc :: Int -> [Int] -> Array Int Int
mkAcc n xs = listArray (0,n) $ scanl (+) 0 xs

公式解説も同じことを言っている、けどちょっとおかしい。これでは割引きなしの定価購入になってしまう。
「商品 $i$ のセール価格を $B_i (= P_i / 2)$ とおくと」とするか、
「かかった金額は $(B_{L_k} + \dots + B_{R_k})/2$ である」とプログラムに合わせるか、
どちらかの修正が必要だ。

小課題3

$M \leq 10$

…これは何を誘っているんだろう？
それぞれの日について累積和を取り直して、日ごとに客をさばく感じかな？

客の順番通りでなくなったので、日ごとに、累積和を作り直して、その日に来た客の答えを配列にため込む。

compute [n,m,q] pas tlrs | m <= 10 =
  elems $ array (1, q)
    [ (k, (pacc ! r) - (pacc ! pred l))
    | d <- [1 .. m]
    , let pacc = mkAcc n [if d == a then div p 2 else p | p:a:_ <- pas]
    , (k, t:l:r:_) <- zip [1 ..] tlrs
    , t == d ]

あるいは、日ごとの累積和を$M$日分全部作っておけば、人ごとに計算できる。

compute [n,m,q] pas tlrs | m <= 10 =
  [ (pacc ! r) - (pacc ! pred l)
  | t:l:r:_ <- tlrs, let pacc = paccs ! t ]
  where
    paccs = listArray (1, m)
      [ mkAcc n [if d == a then div p 2 else p | p:a:_ <- pas]
      | d <- [1..m] ]

公式は後者だった。

小課題4

$A_i \neq A_j \ (i \neq j)$

$A_i$ に重複がない設定。
これは、特定の日に割引きされる商品はたかだか一つということ。
すると、人ごとに、その日に値引きされる商品番号（または値引きなし）が一つだけ決まるので、それを買うかどうかも $O(1)$ で判定できるから、累積和は定価のものだけ用意すれば事足りる。
ただし、日から値引きされる商品番号は引けるように表を作っておく必要がある。ついでに値引き額も付けておく。

compute [n,m,q] pas tlrs | (case4) =
  [ (pacc ! r) - (pacc ! pred l) - if l <= i && i <= r then disc else 0
  | t:l:r:_ <- tlrs
  , let (i, disc) = d2i ! t ]
  where
    pacc = mkAcc n $ map head pas
    d2i = accumArray (flip const) (0,0) (1,m) -- 日付 → (その日の値引き商品番号, 値引き額)
      [ (a, (i, div p 2))
      | (i, p:a:_) <- zip [1..] pas ]

やる意味は無いけど小問題4の条件に合うか判定する計算

    case4 = all (2 >) $ elems $ accumArray (+) 0 (1,m) [(a,1) | p:a:_ <- pas]

小課題5

$P_i = 2$

これはまた奇抜な。商品の値段は全て2円。なので定価購入価格は品数の倍で $2(R-L+1)$ 円。
ここから、日 $T_k$ の値引き商品群のうち、LからRの範囲内の個数ぶんだけ値引きすればいい。

（ちょっと方針に自信が持てなくて解説をカンニング…）

小課題3同様に日ごとに処理する。
その日に値引きされる商品の番号を昇順に格納した配列をつくり、
これをLやRで二分探索すると添え字の差が個数になるという手法。

抽象化すると、値引き商品の番号の集合を作って、L以上R以下の部分集合の要素数を数えている。

これを集合でするとき、Data.IntSetは size が $O(N)$ なので役に立たない。
Data.Set は $O(\log N)$ なので間に合う上に、findIndex で、配列の添え字に相当する情報が同じ時間で取り出せるのでこれを使えばよい。

import qualified Data.Set as S

compute [n,m,q] pas tlrs | (case5) =
  [ (r - pred l) * 2 - disc
  | t:l:r:_ <- tlrs
  , let s = d2is ! t
  , let disc = getIndex (succ r) s - getIndex l s ]
  where
    d2is = fmap S.fromDistinctDescList $ -- 日aに値引きされるiの集合
           accumArray (flip (:)) [0] (1,m) -- 0 は番兵
           [(a, i) | (i, p:a:_) <- zip [1..] pas]
-- 集合のxより小さい最大要素の背番号
    getIndex x s = case S.lookupLT x s of Just y -> S.findIndex y s

あるいは、IntMap を使って、背番号を手で振ることもできる。

import qualified Data.IntMap as IM

compute [n,m,q] pas tlrs | (case5) =
-- ここは上と同じ
  where
    d2is = fmap mkIndexMap $ -- 日aに値引きされるiの背番号map
           accumArray (flip (:)) [0] (1,m) -- 0 は番兵
           [(a, i) | (i, p:a:_) <- zip [1..] pas]
    mkIndexMap xs = snd $ IM.mapAccum incr 0 $ IM.fromList [(x, ()) | x <- xs]
    incr a _ = (a1, a1) where a1 = succ a
-- 集合のxより小さい最大要素の背番号
    getIndex x im = case IM.lookupLT x im of Just (_, k) -> k

この mkIndexMap は xs に制約を付けないが、ここでは、重複のない降順リストとわかっているので、より簡便に、先に番号を振ってしまうこともできる。

    d2is = fmap mkIndexMap $ -- 日aに値引きされるiの背番号map
           accumArray (flip (:)) [] (1, m) -- 番兵は後で入れる
           [(a, i) | (i, p:a:_) <- zip [1..] pas]
    mkIndexMap = IM.fromList . foldr insrt [(0,0)]
    insrt x kvs@((_,v):_) = (x, succ v) : kvs

(真面目に背番号を0から数える必要もないが、次への伏線。）

無意味な条件確認。

    case5 = all ((2 ==) . head) pas

小課題6

追加の制約なし。

公式解説

今までの誘導の集大成。

小課題5では、L から R の間の種類 d な商品の個数がそのまま値引き額だった。
これを、「その日に値引きな商品一覧の番号」と見て、そのような商品だけの価格の累積和も作っておいて、引き算一発で求める。

小課題5の最後のコードで insrt において、「（今日値引きになっている）一つ前の商品の背番号のsuccが自分の背番号」としていたところを「一つ前までの商品までの値引き額の累積和＋自分の値引き額」に変えるだけでほぼ完成。

compute [n,m,q] pas tlrs =
  [ (pacc ! r) - (pacc ! pred l) - (aR - aL)
  | t:l:r:_ <- tlrs
  , let dm = d2ds ! t
  , let Just (_, aR) = IM.lookupLE r dm
  , let Just (_, aL) = IM.lookupLT l dm ]
  where
    pacc = mkAcc n $ map head pas                -- 定価の累積和
    d2ds = fmap mkDescMap $  -- 日aに割引きな商品番号から割引額の累積和マップ
           accumArray (flip (:)) [] (1,m)
           [(a, (i, div p 2)) | (i, p:a:_) <- zip [1..] pas]
    mkDescMap = IM.fromList . foldr step [(0, 0)]
    step (i,x) jas@((_,a):_) = (i, x + a) : jas

この最終的な解答への誘導が見事。
そしてこの解法は、immutableな計算だけ（配列構築は見逃して）で書けるのもうれしい。
〈自分で思いつけなかったのはくやしい、というか、こんなに日付ごとに値引き商品だけの累積和表を作り直して間に合うとは思わなかった。）

自分の解法

AtCoderにある「別解 by Mitsubachi」と同じアプローチだった。

• 日ごとに、何番の商品がいくら値引きされるのかをまとめておく。
• 商品番号を添え字、定価を値、演算は (+) なセグメント木を作る。
• 日ごとに以下を行う
  • その日の値引き商品の値段をセグメント木で下げる
  • その日の客について、セグメント木のクエリで答えを得る
  • セグメント木を元に戻す

効率的なセグメント木は mutable array を必要とするので、命令型なコードになってしまうのが残念。

compute :: [Int] -> [[Int]] -> [[Int]] -> [Int]
compute [_n,m,q] pas tlrs = runST $
  do
    ans <- MUV.new q -- 客ごとの購入額、答えを入れる配列
    st <- makeSegTree (+) 0 $ 0 : map head pas -- 商品idからpriceの和のセグメント木
    forM_ (assocs t2jlr) (\(t,jlrs) -> do  -- 客の来た日tについて
      let ips = a2is ! t -- 今日tの半額商品リスト
      forM_ ips (\(i, p) -> updateSegTree st i (div p 2)) -- 全て半額にする
      forM_ jlrs (\(j, _t:l:r:_) -> do -- 客jの購入額をそれぞれ求める
        v <- querySegTree st l (succ r)
        MUV.write ans j v
        )
      forM_ ips (\(i, p) -> updateSegTree st i p) -- 割引きを戻す
      )
    forM [0 .. pred q] (MUV.read ans)
  where
    t2jlr = accumArray (flip (:)) [] (1,m) [(t,(j, tlr)) | (j, tlr@(t:_)) <- zip [0..] tlrs] -- 日ごとの客リスト
    a2is  = accumArray (flip (:)) [] (1,m) [(a,(i,p)) | (i,p:a:_) <- zip [1..] pas] -- 日ごとの半額商品idとpriceリスト

-- セグメント木の実装は省略

ユーザ解説「より高速な解法」

公式解説と同様の考察により，番号が $L_k$ 以上 $R_k$ 以下で種類が $T_k$ である商品の定価の総和を求めればよいです．

はい。

クエリを先読みし，各種類の商品の定価の総和を管理しながら番号の昇順に平面走査すると，
種類 $T_k$ についての値を $L_k, R_k$ まで操作したタイミングで確認すれば求められます．

ばなな

実装例を読んだら理解できた。ということで本題の「平面走査」について。

まず、素直な表を考える。

• 表の縦軸を営業日
• 表の横軸を商品番号
• 表の項目は、その日のその商品の価格
  （ひとつの商品の列について見ると、特売日の1マスだけ半額、他は全て定価が書かれる。）

人 $i$ の購入価格は、この表の $T_i$ 行の $L_i$ 列から $R_i$ 列までのマスの総和。
この足し算を高速化するため、商品番号について累積和をとった表を考える。
すると、人 $i$ の購入価格は、累積表の $T_i$ 行の $R_i$ 列 － $(L_i-1)$列、で計算できる。

さてここで、この累積を本気で足し算で実行すると、$O(MN)$ かかってしまう。
ひとつの商品が値引きされるのは1日だけなので、累積和の内容は全ての行でほとんどの列で定価の累積和となり、食い違うタイミングは $N$ 箇所しかない。
そこで、元の表を少し改造する。

• 表の縦軸を営業日、追加のもう一行（第0行）
• 表の横軸を商品番号
• 表の項目は、その日のその商品の値引き額、第0行は定価
  （ひとつの商品の列について見ると、特売日の1マスだけ半額、他は全て0が書かれる。）
• 定価での $L$ から $R$ の購入価格は、この表の 0 行の $L$ 列から $R$ 列までのマスの総和
• 人 $i$ の総割引額は、この表の $T_i$ 行の $L_i$ 列から $R_i$ 列までのマスの総和

そして商品番号について累積した表では、

• 定価での $L$ から $R$ の購入価格は、累積表の 0 行の $R$ 列 － $(L-1)$ 列
• 人 $i$ の総割引額は、累積表の $T_i$ 行の $R_i$ 列 － $(L_i-1)$ 列

の4箇所を読めば答えが得られる。

ここでさらに、この表を二次元配列に展開してしまうと、空間がO(MN)かかってしまう。
その代わりに、縦一列、とある商品番号の商品までの累積和、という$M+1$ 個の要素だけを保持する。
すると、列 $i-1$ から次の列 $i$ に移るときに必要な計算は

• 定価の第0行に $P_i$ を足し込む
• 第 $A_i$ 行に $P_i / 2$ を足し込む

の2つだけで終わる。
表全体をメモリに保持しなくなったので、人ひとりにつきこの表の4箇所から情報を読み取る必要があるため、その位置を列ごとに、また、それを誰の結果に足し込む（または引き去る）かという情報とともにあらかじめ集めておき、列を走査するごとに必要なものを取り出す。

走査が終了すると、全ての人の結果も揃っている、という仕組み。

二次元の表を走査する向きを転置すること、また走査の途中で、必要な値が得られた時点で取り出していくこと、がこのアイデアのミソだろう。

さて、これをHaskellで実装するにあたって、$M+1$要素の配列の2箇所を毎回変更するところがつらい。Data.Arrayではこれは$O(M)$かかってしまう。

効率を無視して書くとこうなる。

compute :: [Int] -> [[Int]] -> [[Int]] -> [Int]
compute [n,m,q] pas tlrs = elems ansM
  where
    gather = accumArray (flip (:)) [] (0,n) $ concat $
      [ [(r, (0, j,  1)), (pred l, (0, j, -1))
        ,(r, (t, j, -1)), (pred l, (t, j,  1))]
      | (j, t:l:r:_) <- zip [1..] tlrs ]
    acc0 = listArray (0,m) $ repeat 0
    ans0 = listArray (1,q) $ repeat 0
    (_, ansM) = foldl step (acc0, ans0) $ zip pas $ tail $ elems gather
    step (acc, ans) (p:a:_, tjss) = (acc1, ans1)
      where
        acc1 = accum (+) acc [(0, p), (a, div p 2)]
        ans1 = accum (+) ans [(j, s * acc1 ! t) | (t,j,s) <- tjss]

Data.IntMapで配列を模倣すると$O(\log M)$かかってしまう。(一応ACするが、他の解法より遅い。)
Data.Vector.modify を in-place にさせることもうまくできず、Data.Vector.Mutable を使うしかなかった。

import qualified Data.Vector.Unboxed.Mutable as MUV
import Control.Monad.ST

compute :: [Int] -> [[Int]] -> [[Int]] -> [Int]
compute [n,m,q] pas tlrs = runST $
  do
    acc <- MUV.replicate (succ m) 0
    ans <- MUV.replicate (succ q) 0
    forM_ (zip pas $ tail $ elems gather) (\(p:a:_, tjss) -> do
      MUV.modify acc (p +) 0
      MUV.modify acc (div p 2 +) a
      forM_ tjss (\(t,j,s) -> do
        acct <- MUV.read acc t
        MUV.modify ans ((s * acct) +) j
        )
      )
    forM [1 .. q] (MUV.read ans)
  where
    gather = accumArray (flip (:)) [] (0,n) $ concat $
      [ [(r, (0, j,  1)), (pred l, (0, j, -1))
        ,(r, (t, j, -1)), (pred l, (t, j,  1))] -- 列番号 → 日付、人、±1
      | (j, t:l:r:_) <- zip [1..] tlrs ]

解説のC++コードでは、定価の累積和はaccとは別の変数で管理していたりする。
考えてみれば、無理に平面走査と同時にやる必要もなかったかも。