yukicoder No.1067 #いろいろな色 / Red and Blue and more various colors (Middle)

お題箱より。

解法

箱を並び替えても答えが変わることはないので、 $A_{i}$ の昇順に並べておきます。

各箱から玉を1個ずつ取り出し、ある1つの色 $c$ の玉がちょうど $p$ 個選ばれる場合の数が、どのように計算できるか図示してみましょう。

f:id:betrue12:20200811140555p:plain

$A_{i}$ の昇順に並べたことによって、箱は色 $c$ の玉を含まないゾーンと含むゾーンに分かれます。含まないゾーンの取り出し方は $A_{i}$ の積であり、色 $c$ の玉はもちろん $0$ 個です。

色 $c$ の玉を含むゾーンの取り出し方は以下のDPによって求めることができます。

$dp\lbrack i\rbrack\lbrack j\rbrack =$ 箱 $i, i+1, ..., N-1$ 全てに色 $c$ の玉が1個ずつ含まれているとき、これらの箱から玉を1個ずつ取り出し、色 $c$ の玉がちょうど $j$ 個選ばれる場合の数。

この値は、箱 $i, i+1, ..., N-1$ の全てに含まれている色であればどの色でも同じ値になります。つまりクエリで必要な色 $c$ 全てについて別々に計算する必要がありません。

初期条件は $dp\lbrack N\rbrack\lbrack 0\rbrack = 1$ で、後ろから求めていきます。遷移は

$dp\lbrack i+1\rbrack\lbrack j\rbrack$ を遷移元とし、箱 $i$ から色 $c$ の玉を選ぶ。係数 $1$ で $dp\lbrack i\rbrack\lbrack j+1\rbrack$ に遷移。
$dp\lbrack i+1\rbrack\lbrack j\rbrack$ を遷移元とし、箱 $i$ から色 $c$ でない玉を選ぶ。係数 $A_{i}-1$ で $dp\lbrack i\rbrack\lbrack j\rbrack$ に遷移。

となります。DPパートの計算量は $O(N^{2})$ です。

クエリに答えるには「ある1つの色 $c$ の玉がちょうど $p$ 個選ばれる場合の数」をたくさんの $(c, p)$ について求める必要があります。これは以下のように計算することができます。

まず、箱が色 $c$ を含まないゾーンと含むゾーンの境界を二分探索で探す。色 $c$ を含む箱の最小のインデックスを $x$ とする（存在しない場合は $x=N$ ）。
色 $c$ を含まないゾーンの選び方は $A_{0}\times A_{1}\times\cdots\times A_{x-1}$ なので、事前に累積「積」を計算しておいてそれを参照する。
色 $c$ を含むゾーンの選び方はDPテーブルを参照する。具体的には $dp\lbrack x\rbrack\lbrack p\rbrack$ である。