Codeforces Round #564 C. Nauuo and Pictures

問題概要

$n$ 枚の画像がランダムに表示されるWebサイトがある。各画像には重み $w_{i}$ が設定されていて、画像 $i$ の表示確率は $\frac{w_{i}}{\sum_{j} w_{j}}$ である。

画像を $m$ 回表示させる。このとき画像が1回表示されるごとにその画像の重みが変動する。各画像には0または1の値 $a_{i}$ が定められていて、表示された画像 $i$ について $a_{i} = 0$ であれば重みが-1、 $a_{i} = 1$ であれば重みが+1される。

画像を $m$ 回表示させた後の各画像の重みの期待値を、有理数 $\bmod 998244353$ で求めよ。

C1

制約

$1 \le n \le 50$
$1 \le m \le 50$
$a_{i}$ は0または1
$1 \le w_{i} \le 50$

解法

制約が小さいので、 $n$ 枚の画像それぞれについて別々に計算することを考えてみましょう。答えを求めたい画像を $x$ とします。

そうすると、 $x$ 以外の画像について具体的に何が選ばれたかはあまり興味がありません。

画像 $x$
画像 $x$ 以外の $a_{i} = 0$ である画像たち（グループ0とします）
画像 $x$ 以外の $a_{i} = 1$ である画像たち（グループ1とします）

の3つだけを区別すれば十分です。ここで、グループ0, 1に属する画像の初期重み合計をあらかじめ計算し、それぞれ $s_{0}, s_{1}$ としておきます。

以下のような状態を持つDPを考えます。

$dp\lbrack i \rbrack\lbrack j \rbrack\lbrack k \rbrack =$ 画像をこれまで合計 $i$ 回表示させて、画像 $x$ が $j$ 回、グループ0の画像たちが合計 $k$ 回表示されている確率

グループ1の画像たちが選ばれた回数が入っていませんが、合計回数が分かっているので $i-j-k$ で求められます。これらの情報があれば、この時点での

$x$ の重み $= w_{x} \pm j$ （ $a_{x} = 0$ ならマイナス、 $a_{x} = 1$ ならプラス）
グループ0の画像たちの重み合計 $= s_{0} - k$
グループ1の画像たちの重み合計 $= s_{1} + (i-j-k)$

が求められるので、そこから遷移確率を計算することができます。最終的に $dp\lbrack M \rbrack\lbrack j \rbrack\lbrack k \rbrack$ を使えば期待値を計算できます。

ACコード

Submission #55259545 - Codeforces

重みが負になるケースを無理やり計算しようとするとおかしなことが起こり得るので注意。負になるケースに正の確率で遷移することはないので、 if(dp[i][j][k] == 0) continue; のような1行を入れておけば大丈夫です。

C2

制約

$1 \le n \le 2\times 10^{5}$
$1 \le m \le 3000$
$a_{i}$ は0または1
$1 \le w_{i}$
$\sum_{i} w_{i} + m \lt 998244353$

解法

C2では制約が大きいので全部別々に計算することはできません。

$a_{i} = 0$ である画像たちをグループ0、 $a_{i} = 1$ である画像たちをグループ1とします。それぞれの初期重みの和を $s_{0}, s_{1}$ とします。

まずは1つ1つの画像が何回選ばれるかを考えずに、グループ0, 1それぞれの合計表示回数が確率いくらで何回になるかを考えます。すなわち以下のようなDPを考えます。

$dp\lbrack i \rbrack\lbrack j \rbrack =$ 画像をこれまで合計 $i$ 回表示させて、グループ0の画像たちが合計 $j$ 回表示されている確率

これはC1と同じような遷移で、状態数 $O(m^{2})$ で計算できます。

ここから1つ1つの画像について考えていきます。ある1つの画像 $x$ について考え、例えばそれがグループ0に属しているとします。グループ0の画像が $k$ 回選ばれた時の画像 $x$ の重みの期待値を $E_{k}$ と表記すると、画像 $x$ についての答えは $\sum_{k} (E_{k} \times dp\lbrack M \rbrack\lbrack k \rbrack)$ となります。

初期値は $E_{0} = w_{x}$ です。 $E_{k}$ から $E_{k+1}$ を計算できないか考えてみましょう。

グループ0の画像を $k$ 回選んでいる時点で、グループ0の重み合計は $s_{0}-k$ です。そのうち画像 $x$ の重みがいくらかによって遷移確率が決まり、その値は確率によって変動しますが…その期待値はまさしく $E_{k}$ として求められています。そのため $k+1$ 回目で画像 $x$ が選ばれる確率は $\frac{E_{k}}{s_{0}-k}$ とみなすことができて、その確率で $x$ の重みが1減るため

$E_{k+1} = E_{k} - \frac{E_{k}}{s_{0}-k} = E_{k}(1 - \frac{1}{s_{0}-k})$

と計算することができます。

さらにこの式から、 $E_{k}$ の値は初期値 $E_{0} = w_{x}$ に比例していることが分かります。つまり全ての $x$ についてこの $E_{k}$ の値を全部計算しなくても、 $E_{0} = 1$ とみなして $E_{k}$ を求め、 $\sum_{k} (E_{k} \times dp\lbrack M \rbrack\lbrack k \rbrack)$ まで計算してしまってから、最後に各 $x$ ごとに $w_{x}$ を掛けることで答えを求めることができます。