ARMERIA

Rubyと競技プログラミングの話 AtCoderやCodeforcesの問題解説記事が多め。

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

お題箱より。

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

解法

箱を並び替えても答えが変わることはないので、 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) について求める必要があります。これは以下のように計算することができます。

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

これは1つの  (c, p) あたり  O(\log N) で計算できます。必要な  (c, p) の個数は  \sum_{i=1}^{Q}(r_{i}-l_{i}+1) \le 201\times 5000 なので、ちょっと厳しいですが間に合います。

ACコード

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

公式解説では多項式の係数として解説していますが、1次式を掛けた時の係数の計算と今回のDPは全く同じものなので、解法としては同じです。