Educational Codeforces Round 11 E. Different Subsets For All Tuples

お題箱より。

Problem - E - Codeforces

公式解説より楽な解法があるのでそっちを紹介します。記事最後に、リクエストにあった公式解説の補足についても説明します。

問題概要

整数列 $a$ に対して、 $f(a)$ を「 $a$ の部分列（連続でなくても良い）として登場する長さ $0$ 以上の相異なる数列の個数」と定義する。

各要素が $1, ..., m$ のいずれかであるような長さ $n$ の数列は $m^{n}$ 個ある。その全てについて $f(a)$ の値を合計した値を $\bmod 10^{9}+7$ で求めよ。

制約

$1 \le n, m \le 10^{6}$

解法

※解説中の「数列」という言葉は、特に断らない限り各要素が $1, ..., m$ のいずれかであるようなものを指すことにします。

数え方についての考察

$f(a)$ の計算ルールとして、相異なる部分列の個数を数えるので、例えば $\lbrack 1, 2, 1, 3, 1, 2, 3\rbrack \to \lbrack 1, 2, 3\rbrack$ のように複数の取り出し方で同じ部分列ができる場合にも $1$ 通りとして数える必要があります。このときによく使うのが、「その部分列を取り出す位置の組み合わせのうち、最も左側から取り出すようなものだけを数える対象とする」、という考え方です。先ほどの例の場合は（先頭を $0$ 番目として） $0, 1, 3$ 番目が最も左側から取り出す位置になります。

これを踏まえて、ある長さ $k$ の数列 $s$ があった時に、それを部分列として含む長さ $n$ の数列 $a$ を数える方法を考えます。まず、先ほど考えた「最も左側から取り出す」取り出し方で $s$ の各要素に対応する $a$ の位置を決めたとします。そうすると $a$ の各要素の値の候補数は、

$s$ の要素に対応している： $1$ 通り
$s$ の要素に対応していない＆それ以降に $s$ に対応する要素がまだ存在する： $m-1$ 通り
$s$ の要素に対応している＆それ以降に $s$ に対応する要素がもう存在しない： $m$ 通り

となります。 $2$ 番目が $m-1$ 通りとなる理由は、もし $s$ における次の要素と同じ値を採用してしまった場合、そこから取り出したほうがより左側になってしまうので、「最初に決めた位置が最も左側から取り出す位置である」という決め事に反するからです。

この $a$ が何通りあるかは、 $s$ の要素に対応する最後の位置が何番目であるかによって変わってきます。公式解説ではこの「 $s$ の長さはいくつか」「 $s$ の要素に対応する最後の要素は $a$ の何番目か」を組み合わせた2重和を考えていますが、別の方法があります。DPをしましょう。

DPをする

取り出し元の数列 $a$ と部分列として取り出す数列 $s$ の組 $(a, s)$ をDPで数えます。 $a$ の要素数を $1$ つずつ増やしていきながら、 $s$ の要素数を増やすパターン、増やさないパターンのそれぞれの遷移をします。

先ほど見たように、 $s$ の要素に対応していない $a$ の要素の候補数は、まだ $s$ に対応する要素が残っているかどうかで変化します。なのでこれを考慮し、次のような状態定義をします。

$dp\lbrack i\rbrack\lbrack 0\rbrack =$ 長さ $i$ の数列 $a$ と、その部分列である長さ $0$ 以上の数列 $s$ の組であって、今後 $s$ の要素を増やす予定のもの
$dp\lbrack i\rbrack\lbrack 1\rbrack =$ 長さ $i$ の数列 $a$ と、その部分列である長さ $0$ 以上の数列 $s$ の組であって、もう $s$ の要素を増やさないもの

その時点での $s$ の長さは持たなくても良い、という点が1つの注目ポイントです。

$dp\lbrack i\rbrack$ からの遷移は以下のように列挙できます。

$dp\lbrack i\rbrack\lbrack 0\rbrack$ から、 $a$ の要素を末尾に $1$ つ増やし、 $s$ の要素としては採用しない： $m-1$ を掛けて $dp\lbrack i+1\rbrack\lbrack 0\rbrack$ に遷移
$dp\lbrack i\rbrack\lbrack 0\rbrack$ から、 $a$ の要素を末尾に $1$ つ増やし、それを $s$ の要素としても末尾に追加する： $m$ を掛けて、 $dp\lbrack i+1\rbrack\lbrack 0\rbrack$ （まだ終わらない）と $dp\lbrack i+1\rbrack\lbrack 1\rbrack$ （これで終わり）の両方に遷移
$dp\lbrack i\rbrack\lbrack 1\rbrack$ から、 $a$ の要素を末尾に $1$ つ増やし、 $s$ の要素としては採用しない： $m$ を掛けて $dp\lbrack i+1\rbrack\lbrack 1\rbrack$ に遷移

初期条件は $dp\lbrack 0\rbrack\lbrack 0\rbrack = dp\lbrack 0\rbrack\lbrack 1\rbrack = 0$ とすれば良いです。後者は $s$ として何もない状態からもう要素を増やさない、つまり空数列に対応します。答えは $dp\lbrack n\rbrack\lbrack 1\rbrack$ になります。