概要

「条件を満たすものすべてについてのスコアの総和を求めよ」という形の問題を DP で解けることがある。このときスコアの総和だけをもってもうまくいかなくて、条件を満たすものの個数も同時にもっておく必要がある。

例題

$N$ 個の整数 $A_1, A_2, \ldots, A_N$ がある。隣り合う $2$ 数の間に + または - を入れてできる式のうち、- が $2$ 回以上連続で登場しないものを「良い式」とよぶ。すべての良い式の値の総和を $\bmod 10^9 + 7$ で求めよ。 $N \leq 10^5$

良い式の個数を数える問題

まず良い式の個数を数える問題を考えます*1。前から + と - のどちらを入れるか決めていくときに、「直前が - だったら - は使えない」というルールなので、「どこまで決めたか」「直前が + か - か」を状態としてもてば DP できそうです。つまり

状態 $(i, op) :=$ $A_i$ の直前までの入れ方を決めたとき、最後に入れた演算子が $op$ であるような状態
$dpc[i][j] :=$ $j = 0$ なら $op =$ +、 $j = 1$ なら $op =$ - として、状態 $(i, op)$ にする方法の個数*2

とできます。あとは「直前が - なら - は使えない」のルールを考慮して、状態 $(i,$ + $)$ からは状態 $(i + 1,$ + $)$ または状態 $(i + 1,$ - $)$ 、状態 $(i,$ - $)$ からは状態 $(i + 1,$ + $)$ に遷移できることにすればよいです。ということで遷移は

$dpc[i + 1][0] = dpc[i][0] + dpc[i][1]$
$dpc[i + 1][1] = dpc[i][0]$

と書けました。

良い式の総和を求める問題*3

式の値というのは、+ を入れたらスコアに $A_i$ が、- を入れたらスコアに $-A_i$ が加算される、というふうにしたときのスコアだと考えることができます。

個数の問題の遷移をほぼそのまま総和の問題にも適用して

$dps[i][j] :=$ $j = 0$ なら $op =$ +、 $j = 1$ なら $op =$ - として、状態 $(i, op)$ にする方法すべてについてのスコアの総和*4
（間違い） $dps[i + 1][0] = (dps[i][0] + A_{i + 1}) + (dps[i][1] + A_{i + 1})$
（間違い） $dps[i + 1][1] = (dps[i][0] - A_{i + 1})$

としたくなります*5が、これはうまくいかないです。なぜでしょう。

一言でいうと、「たくさんある式を状態にまとめており、本来式の数ぶん足さないといけないのに、状態の数ぶんしか足していないことになるから」です。

たとえば $N = 4, (A_1, A_2, A_3, A_4) = (2, 7, 1, 8)$ のケースで、状態 $(2,$ + $)$ 、つまり $2 \ ? \ 7 \ ? \ 1$ の部分まで決まっていて、最後が + なので $2 \ ? \ 7 + 1$ となっている状態を考えてみましょう。この状態になる方法は $2 + 7 + 1$ と $2 - 7 + 1$ の $2$ つあるので、ここから + を入れることにして状態 $(3,$ + $)$ に遷移したとすれば $2 + 7 + 1 + 8$ と $2 - 7 + 1 + 8$ の $2$ つを考える必要が出てきます。このとき $8$ は $2$ 回足されていないといけないですね。それなのに、上の間違った遷移では $8$ は $1$ 回しか足されないことになってしまいます。

そこでどうするかというと、良い式の個数も同時に考えればよいです。上の例で $8$ を $2$ 回足すべきだというのは、状態 $(3,$ + $)$ になる方法の数が $2$ つあるからです。つまり、スコアへの寄与をありうるものの個数ぶん足せばよいわけです。ということで

$dps[i + 1][0] = (dps[i][0] + dpc[i][0] \cdot A_{i + 1}) + (dps[i][1] + dpc[i][1] \cdot A_{i + 1})$
$dps[i + 1][1] = (dps[i][0] - dpc[i][0] \cdot A_{i + 1})$

と書けました。

実装例

C++ with ACL です（遷移を読みやすくするために modint を使いました）。

#include <bits/stdc++.h>
using namespace std;
#include <atcoder/modint>
using namespace atcoder;
using mint = modint1000000007;

int main()
{
  int N;
  cin >> N;
  vector<int> A(N);
  for (int i = 0; i < N; i++)
  {
    cin >> A[i];
  }

  vector<vector<mint>> dpc(N, vector<mint>(2, 0));
  vector<vector<mint>> dps(N, vector<mint>(2, 0));
  dpc[0][0] = 1; // 最初は - が使えるので、直前が + だったとみることにする
  dps[0][0] = A[0];
  for (int i = 0; i < N - 1; i++)
  {
    // 遷移を直接書いてもよいが、ここではループで書く方法を紹介する
    for (int j = 0; j < 2; j++)
    {
      for (int nj = 0; nj < 2; nj++) // nj は遷移先の j
      {
        if (j == 1 && nj == 1) // - は 2 回連続で使えない
          continue;

        mint a = (nj == 0 ? A[i + 1] : -A[i + 1]); // スコアへの寄与
        dpc[i + 1][nj] += dpc[i][j];
        dps[i + 1][nj] += dps[i][j] + dpc[i][j] * a;
      }
    }
  }
  mint ans = dps[N - 1][0] + dps[N - 1][1];
  cout << ans.val() << endl;
}

期待値の線形性との関連

$dpe[i][j] := \frac{dps[i][j]}{dpc[i][j]}$ *6と定義して、上の遷移を変形してみましょう。すると

$dpe[i + 1][0] = \frac{dpc[i][0]}{dpc[i][0] + dpc[i][1]} dpe[i][0] + \frac{dpc[i][1]}{dpc[i][0] + dpc[i][1]} dpe[i][1] + A_{i + 1}$
$dpe[i + 1][1] = dpe[i][0] - A_{i + 1}$

が得られます（難しい変形ではないはずなので、気になる人は各自やってみましょう）。この式、なんだか「期待値の線形性」っぽいですよね。下の式は本当にそのままな感じがしますし、上の式も $\frac{dpc[i][0]}{dpc[i][0] + dpc[i][1]}, \frac{dpc[i][1]}{dpc[i][0] + dpc[i][1]}$ を、状態 $(i + 1,$ + $)$ にくる前にそれぞれ状態 $(i,$ + $)$ 、状態 $(i,$ - $)$ であった確率、と捉えることができそうです。

類題

（このテク単体で解ける問題が探せなかった……。知ってる人いたら教えてください）

ほかのテクと組み合わせるやつ：

まとめ

結局、ありうるものの個数をもっておく必要があったのは、スコアへの寄与が足される回数がありうるものの個数になるからでした。慣れてないとけっこう混乱しそうなので、この $dps[\mathrm{state_{new}}] \ \xleftarrow{+} \ dps[\mathrm{state_{old}}] + dpc[\mathrm{state_{old}}] \cdot (スコアへの寄与)$ の形は典型として捉えておくのがよさそうです。

*1:実は "+" か "-+" を使うと言い換えるとフィボナッチ数であることがわかるんですが、ここではより問題設定に忠実にやります

*2:dpc の c は count の c です

*3:こちらも問題の性質が良すぎるためにもっと楽な見方があります、気になる方は公式解説を見てみましょう。これから説明する解法はより汎用的だと思っています

*4:dps の s は sum の s です

*5:別にならないが？って人もいるかも

*6:dpe の e は期待値の英語 expected value の e です

miscalc のブログ

主に競プロの話をします

個数と総和をもつ DP

概要

例題

良い式の個数を数える問題

良い式の総和を求める問題*3

実装例

期待値の線形性との関連

類題

まとめ