A - 321-like Checker

問題 ABC321A

「次の数字が、より小さいものである」が全ての数字について満たされていることを確認できればよい。
なので数として読み込まない。

結果

abc321a :: String  -- Nの文字列
        -> Bool    -- 答え
abc321a s = and $ zipWith (>) s $ tail s

B - Cutoff

問題 ABC321B

シグネチャを決める。

abc321b :: Int   -- N
        -> Int   -- X
        -> [Int] -- Ai
        -> Int   -- 答え

（総当たりも二分探索もいらないよね…）

$S = \sum_{i<N} A_i$ ここまでの和
$A_{\min} = \min_{i < N} A_i$ ここまでの最低点
$A_{\max} = \max_{i < N} A_i$ ここまでの最高点
$T = S - A_{\min} - A_{\max}$ 必ず最終点に含まれる点数の和

とおく。
最終回が普通の範囲内の点数をとったとき($A_{\min} < A_N \leq A_{\max}$)、最高点と最低点は据え置きで、最終結果は $T + A_N$ となる。
これが目標点と等しくなると $T + A_N = X$ より $A_N = X - T$ ただし $A_{\min} < A_N \leq A_{\max}$ となる。

実際、$A_{\max} < A_N$ とは最後に最高点を更新した場合となり、無視される最高点が交代し、最終結果は $T + A_{\max}$ となる。これが獲得できる点数の上限で、これを上回る$X$は達成できない。

また $A_N \leq A_{\min}$ とは最後に最低点を更新した（または同着となった）場合となり、無視される最低点が交代し、最終結果は $T + A_{\min}$ となる。これが獲得できる点数の下限で、$X$がこれを下回るとき、$A_N$はいくつでも構わないので0でよい。

結果

abc321b :: Int -> Int -> [Int] -> Int
abc321b _n x as
  | an <= amin = 0
  | amax < an = -1
  | otherwise = an
  where
    amax = maximum as
    amin = minimum as
    an = x - sum as + amax + amin

C - 321-like Searcher

問題 ABC321C

シグネチャを決める。

abc321c :: Int  -- K
        -> Int  -- 答え

$K$の上限が示されていないが、321-like Numberの最大値は 9876543210 で、実際これが1022番。
10桁の数なので、A問題と同じ条件だからといって、総当たりで確認していくnaiveな方法はとれない。（like321sがコンパイル時に計算されれば何とか。）

abc321c :: Int -> Int
abc321c k = like321s !! pred k

like321s :: [Int]
like321s = [i | i <- [1..], is321like i]

is321like :: Int -> Bool
is321like n = and $ zipWith (>) s $ tail s
  where
    s = show n

DPな解法

C問題には大げさだが、思いついたやり方。

最上位桁が $d$ で、$m+1$ 桁の 321-like Number の昇順のリストを配列$arr[m,d]$にとる。
ただし0も含める。
最上位桁が $d$ で 1桁の数は、それぞれ $d$ のみである。$arr[0,d] = [d]$
最上位桁が $d$ で $m$桁の数は、$e = 0,\dots,d-1$ に関する $arr[m-1,e]$ の要素に対して $d$ を最上位桁に追加すればよい。

この$arr$の要素を添え字の辞書順に取り出せば、先頭の0を除いて昇順の321-like Numberの列になる。

abc321c :: Int -> Int
abc321c k = like321s !! k

like321s :: [Int]
like321s = concat $ elems like321A

like321A :: Array (Int,Int) [Int]
like321A = listArray bnds $ map f $ range bnds
  where
    bnds = ((0,0),(9,9))
    f (0,d) = [d]
    f (m,d) = [d * 10^m + y | e <- [0..pred d], y <- like321A ! (pred m, e)]

公式解説のやり方

公式解説 by physics0523

それぞれの数字はたかだか1度しか使えないことから、それぞれの数字を使うかどうかで $2^{10}$ の可能性がある。
数字が降順になるという条件から、どの数字を使うかという選択をした時点で、そのような321-like Numberは（たかだか）ひとつしかない。例外は0。

$2^{10}-1$までを数え上げ、ビットの立っている数字を大きい方から並べた321-like Numberを作り、後で全体を整列させる。

import Data.Bits
import Data.List

bTo321 :: Int -> Int
bTo321 x = foldl (\acc d -> acc * 10 + d) 0 [b | b <- [9,8..0], testBit x b]

like321s = sort $ map bTo321 [1..1023]

ユーザ解説のやり方

ユーザ解説 by cirno3153
なんだか変なJavaだなぁ、何だこの it とか until とか、と思ったらKotlinだそうで。
知らない言語のコードだけではとても読み取れない。PlayGroundでいじくって、何をしているのかを推測した。

321-like Number に対して、その続きに（最下位桁より小さい）数字をもう一つ付け加えても321-like Numberになる。

という性質を使って、最初から昇順に生成させている。なるほどねぇ。

like321s = [1..9] ++ concatMap extend like321s

extend :: Int -> [Int]
extend x = [x * 10 + d | d <- [0..pred $ mod x 10]]

0を含めていないので他の定義とは1ずれる。
この定義は停止させていないので止まらないが、正しい$K$だけが与えられるなら問題ない。

D - Set Menu

問題 ABC321D

シグネチャを決める。

abc321d :: Int  -- N
        -> Int  -- M
        -> Int  -- P
        -> [Int]  -- Ai
        -> [Int]  -- Bi
        -> Int  -- 答え

$A_i + B_j$ の総和を求めるが、上限 $P$ でクランプした値で考えることが求められている。
$A_i$ に対して、$B_j$ がかなり高額、具体的には $P - A_i < B_j$ ならばクランプされ、そうでないなら $B_j$ の実際の値を用いる。
つまり、$A_i$ に対して、$P - A_i$ 未満であるような $B_j$ の個数 $cnt$ とその総和 $acc$ が得られれば、$cnt$ 個の組み合わせはそのまま、$M - cnt$ 個の組み合わせはクランプされるので、合計は $A_i * cnt + acc + P * (M - cnt)$ となるとわかる。

そういう計算は、競プロ的には二分探索、HaskellではIntMapで解く。

結果

import qualified Data.IntMap as IM

abc321d :: Int -> Int -> Int -> [Int] -> [Int] -> Int
abc321d _n m p as bs = sum
    [ a * cnt + acc + p * (m - cnt)
    | a <- as
    , let Just (_,(cnt,acc)) = IM.lookupLT (p - a) bmap]
  where
    bmap = IM.fromAscList $ scanl step (minBound, (0,0)) $ sort bs
    step (_,(cnt,acc)) b = (b, (succ cnt, acc + b))

E - Complete Binary Tree

問題 ABC321E

シグネチャを決める。

abc321e :: Int  -- N
        -> Int  -- X
        -> Int  -- K
        -> Int  -- 答え

考える

頂点 $X$ から距離 $K$ の頂点は、以下の3種類に分類できる。

頂点 $X$ から葉の方に $K$ 進んだ先の頂点（$K=0$のとき$X$自身の1個）
頂点 $X$ から根の方に $i$ 進み、来た道でない方の子に降りた頂点 $Y$ から葉の方に $K - i - 1$進んだ先の頂点（$1 \leq i < K - 1$）
頂点 $X$ から根の方に $K$ 進んだ先にある先祖1個（$K=0$のとき$X$自身は数えない）

K=0のとき

$K = 0$ の場合は別扱いにしてしまうのが早い。

abc321e _ _ 0 = 1

K>0のとき

上と下の二手に分かれる。子を数える計算は、自分だけでなく兄弟にも適用する。

abc321e n x k = upwards n x k + downwards n x k

下に進む

頂点$a$の左の子は $2a$、右の子は $2a + 1$ なので、$X$ から始めて、$K$ 回この計算を繰り返す。
最後の頂点$N$が$X$の子孫で、$K$がちょうど葉に到達する値のとき、右の子孫は存在しないので、個数を数えるときは右を$N$でクランプする。

また、途中で左が$N$を超えるような場合は、$K$が大きすぎて木をはみ出しているので計算を打ち切る。

downwards :: Int -> Int -> Int -> Int
downwards n x k = loop x x k
  where
    loop l r 0 = max 0 $ min n r - l + 1
    loop l r k
      | n < l = 0
      | otherwise = loop (l + l) (r + succ r) (pred k)

1段階ずつ降りる代わりに、$K$ビットシフトを用いて一気に最終的な l,r を求めることもできるが、その方法はIntをはみ出す危険がある。

上に進む

該当する頂点が存在しない場合への対処を容易にするため、1ステップずつ先祖に進み、まだ進めるかを確認しながら、兄弟の子孫と直接の先祖の両方を数える。

現在位置$X$に対して、親$P = \lfloor X/2 \rfloor$、左の子 $2P$ 右の子 $2P+1$、よって $X$ でない方の兄弟 $Y = 2P + 2P + 1 - X$ である。

upwards :: Int -> Int -> Int -> Int
upwards _n _x 0 = 1 -- 距離0は現在位置そのもののみ
upwards _n 1 _k = 0 -- 根は親を持たないので距離K>0の頂点もない
upwards _n _x 1 = 1 -- 根でない現在位置から距離1の頂点は親のみ
upwards n  x  k =   -- 親に1進んでK-1した結果と、兄弟YからK-2の子孫
    upwards n p (pred k) + downwards n y (k - 2)
  where
    p = div x 2
    y = p * 4 + 1 - x

F - #(subset sum = K) with Add and Erase

問題 ABC321F

シグネチャを決める。

abc321f :: Int  -- Q
        -> Int  -- K
        -> [(Bool, Int)]  -- Queryiが '+' かどうかと、x
        -> [Int]  -- 答え

わからなくて公式解説を見た。
（これは、immutable脳では逆にわからなくなる発想だなぁ。）

+ x のクエリだけを考える。と、初歩的なDPの問題になる。
0～$K$に関して、それらを作るやり方の場合の数を $cnt[0 ～ K]$ とする。
ボールがひとつもない初期状態では $cnt[0] = 1, cnt[i] = 0$ である。
$x$ なボールを追加したとき、$i-x$ を作るやり方にこのボールを追加すると $i$ が作れることから、$cnt'[i] = cnt[i] + cnt[i - x]$ と足し込むことで次の状態が作れる。
（ただし範囲は全域 $0 \leq i \leq K$、はみ出す $i - x < 0$ のとき $cnt[i - x] = 0$ とする。）

ここで、immutable脳では $cnt$ と $cnt'$ は別物なので加算は独立に行える。

cnt' = take x cnt ++ zipWith (+) (drop x cnt) cnt

一方手続き脳では配列を上書きして使い回すために、

for (i = x; i <= K; i++) { cnt[i] += cnt[i - x]; }

と添え字の小さい方から計算すると、例えば $cnt'[2x] = cnt[2x] + cnt[x]$ となるべきところが、先行する代入のために $cnt[2x] + (cnt[x] + cnt[0])$ と違ってしまう。そこでループの向きを変えて

for (i = K; i >= x; i--) { cnt[i] += cnt[i - x]; }

と、もう参照されない、添え字の大きい方から計算することでこの問題を回避する。

次に、- x のクエリに対応することを考える。
状態として保持しているカウント値は「場合の数」なので、+ x が与えられた順序には依存しない。なので、+ xが最後に与えられた結果が現状だと考えて、１つ前の状態を復元することを考える。

しかし、単純にずらして引いても結果は正しくならない。

cnt : 復元したい、元の状態
cnt' : + x した、手元にある情報 $cnt'[i] = cnt[i] + cnt[i-x]$
cnt'' : 復元のために、xずらして引いた結果 $cnt''[i] = cnt'[i] - cnt'[i-x]$

i	cnt	cnt'	cnt''
$0$	$c_0$	$c_0$	$c_0$
$x$	$c_1$	$c_1 + c_0$	$c_1 - c_0$
$2x$	$c_2$	$c_2 + c_1$	$c_2 - c_0$

ここで、先の手続き脳を逆利用して、引く数も保持している配列を破壊的に更新して、「引く数が正しくなったものから順に引いていく」ように、添え字を小さい方から回す。

for (i = x; i <= K; i++) { cnt[i] -= cnt[i - x]; }

i	cnt	cnt'	手続き脳
$0$	$c_0$	$c_0$	$c_0$
$x$	$c_1$	$c_1 + c_0$	$c_1$
$2x$	$c_2$	$c_2 + c_1$	$c_2$

結果、見事に cnt が復元された。あっぱれ。

さてこれを immutable でどう再現するか。$cnt'[2x]$ からは $cnt'[x]$ ではなく計算済みの $cnt''[x]$ を引く必要がある。
こういうときは accumArray は使えず、フィボナッチ数列の計算のような自己参照をして達成する。

cnt'' = take x cnt' ++ zipWith (-) (drop x cnt') cnt''

結果

import Control.DeepSeq

abc321f :: Int -> Int -> [(Bool,Int)] -> [Int]
abc321f _q k sxs = ans
  where
    count0 = 1 : replicate k 0
    ans = map last $ tail $ scanl step count0 sxs
    step count (b, x) = force $ if b then countP else countM
      where
        countP = take x count ++ zipWith add (drop x count) count
        countM = take x count ++ zipWith sub (drop x count) countM

modBase = 998244353
reg x = mod x modBase
add x y = reg $ x + y
sub x y = reg $ x - y