むかでのチラ裏

競技プログラミング現役引退でチラ裏に・・・

AtCoder Regular Contest 107 D - Number of Multisets

もはやなんで苦手なのかわからないが,苦手だったのでまとめた.

特に特殊ではないと思うが,着眼点をミスるとハマってしまいそうなDPの問題でした.

atcoder.jp

問題概要

自然数 $ N, K $ が与えられて,$ \frac{1}{2^{i}} (i = 0, 1, 2, \cdots) $ のみを要素に使って作れる長さ $ N $,和 $ K $ の多重集合 (multiset) の種類数を答えよ. 答えは大きくなるので $ 998244353 $ で割った余りを出力せよ.

制約

  • $ 1 \leq K \leq N \leq 3000 $
  • 入力は整数

考察

問題を読み替える.$ K $ は整数なので $ \frac{1}{2^{N-1}} $ 未満の数字は使うことができない. (少し考えると $ \frac{1}{2} + \frac{1}{4} + \cdots + \frac{1}{2^{N-1}} $ が小さい数を使う限界であるとわかる)

$$ \sum _ {i=0}^{N-1} c_i \cdot \frac{1}{2^{i}}= K \tag{1} $$ $$ \sum _ {i=0}^{N-1} c_i = N \tag{2} $$ $$ c_i \geq 0 ~~ (0 \leq i \leq N-1) \tag{3} $$

をみたすような整数列$ c $ の数を求めよ,という問題に読み替えることができる.

特殊な制約を持った部分和DPのような見た目をしていて,制約もそんなに厳しくないのでDPの方向で考えていく.

何も考えないで部分和DPをやるととんでもない状態数になることは目に見えているのでこの問題特有の制約について考える.

ここで式 $ (1) $ に注目して両辺に $ 2 $ を掛ける.これは高校数学典型っぽい変形テクニックで $$ \sum _ {i=0}^{N-1} C_i \cdot r^{i} = S $$ のような形が出てきたら,両辺に $ r $ か $ r^{-1} $ を掛けることによって変形すると嬉しいことが多い.

\begin{align} \sum _ {i=0}^{N-1} c_i \cdot \frac{1}{2^{i-1}} &= 2K \\ 2c_0 + \sum _ {i=1}^{N-1} c_i \cdot \frac{1}{2^{i-1}} &= 2K \\ \sum _ {i=0}^{N-2} c_{i+1} \cdot \frac{1}{2^{i}} &= 2(K - c_0) \tag{4} \end{align}

また $$ \sum _ {i=0}^{N-2} c_{i+1} = N - c_0 \tag{5} $$

式 $ (4)~(5) $ は式 $ (1)~(2) $ に対応していることがわかる.ここで入力が $ N, K $ の時の答えを $ f(N,K) $ と置き $ c_0 $ は変数であることに注意すると次の式が成立する.

$$ f(N,K) = \sum _ {c_0=0}^{\min({N,K})} f(N - c_0, 2(K - c_0)) \tag{6} $$

$ \min({N,K}) $ の理由は $ N<0 $ または $ K<0 $ のとき $ f(N,K)=0 $ であるためそれぞれが非負の区間で $ c_0 $ を回したということ.

式 $ (6) $ のような遷移式が見つかったのでDPを考える. $ N<K $ のとき明らかに $ f(N,K)=0 $ のためこの状態を考える必要はなく,状態数は $ O(N^{2}) $ である.遷移が $ O(N) $ なので全体の計算量は $ O(N^{3}) $ になる.$ N \leq 3000 $ なのでこのままでは間に合わない.遷移をまとめることを考えていく.

分かり易さのために $ f(N,K) $ の値を表にして,例として $ f(4,2) $ の値を求めることを考える.

f:id:spider_0859:20210104174424p:plain:w300
☆の値は赤い部分の値の和に等しい

式 $ (6) $ より,赤い部分の値の和が目的の $ f(4,2) $ の値と等しいことがわかる.この赤い部分の和を斜めのラインの累積和として求めながらDPをすれば遷移が $ O(1) $ になるだろう.どう累積和をとってもいいが,今回は図の"〇"のようなこれ以上左上に行けない部分を代表としてそこの添え字に値を加算していった.

注意点としては $ f(N,K) $ の値が $ f(N,2K) $ に依存しているため $ K $ の降順に計算していく.

これでも全体の計算量が $ O(N^{2}) $ になるのだが累積和を取るところが少し面倒.実はもっと良い方法がある.

f:id:spider_0859:20210104175834p:plain:w300
★の値は緑の和なので☆=★+赤となる

先程と同様に考えると★の値は緑の部分の値の和になる.これは☆の値を求める赤の和と殆ど被っている.またそのような★の位置は☆の位置のちょうど左上である(冷静に式を処理すると確認できる).以上のことから次の式が成立する.

$$ f(N, K) = f(N - 1, K - 1) + f(N, 2K) \tag{7} $$

式 $ (7) $ からメモ化再帰なりすれば答えが $ O(N^{2}) $ で得られる.ループを回すときはやはり $ K $ の降順に計算していくことに注意.

別解

AtCoderにおける想定解と同じ.

答えは $ c_0 = 0 $ の場合と $ c_0 > 0 $ の場合の和になることに注目する.

  • $ c_0 = 0 $ の場合は式 $ (1) $ の両辺に $ 2 $ を掛けることで $ f(N, 2K) $ と等しいことが示せる.
  • $ c_0 > 0 $ の場合は $ c_0 $ を $ 1 $ 減らす.すると $ f(N - 1, K - 1) $ と等しいことがわかる.

よって式 $ (7) $ が成立することが示せる.

この場合分けを一発でできる人天才すぎないか?僕には無理です.

個人的ポイント

  • 日本語が多い問題文だったので式で問題文を読み替えることが大切
  • 等比数列に何らかの係数がついた式の総和は公比を掛け/割りすると嬉しい
    • 最初の式と同じようなものが得られて一般的な式が得られがち
  • 遷移がまとめられないかを考える時はDPテーブルを実際に書いてみると良いかも
  • 本文には書かなかったがsnukeさんの解説では1の寄与を二分木1つに言い換えていた
    • 二分木は結構嬉しい形なのでこの言い換えは詰まったときに一考の余地があると思った
    • 二分木を数え上げる時は葉から構成するといいらしい(AtCoder解説動画情報)
      • 実際根から素直に構成すると $ O(N^{3}) $ になってしまう.

ソースコード

提出URL1(累積和)

提出URL2(メモ化再帰)

▶ソースコードを展開