AtCoder Grand Contest 046 C - Shift

C - Shift

解法

$S$ に含まれる 0 の個数を $N_{0}$ 、1 の個数を $N_{1}$ と書くことにします。 $S$ を、0 で区切られた $N_{0}+1$ 個のエリアに分割します。

f:id:betrue12:20200621170809p:plain

図に示すように、エリア $i$ に含まれる 1 の個数を $one\lbrack i \rbrack$ とし、エリア $i$ を含むそれ以前に含まれる 1 の個数を $sum\lbrack i \rbrack$ とします。

本問題の操作は、1 を1つ選んで元々あったエリアより左のエリアに移動させる操作だと思うことができます。つまり操作後の状態において、以下の条件が満たされている必要があります。

1 の個数の合計は初期状態と同じである。
初期状態よりも 1 が増加しているエリアについてその増加量を合計した値は $K$ 以下である。
エリア $i$ 以前（エリア $i$ を含む）に存在する 1 の個数は $sum\lbrack i \rbrack$ 個以上である。

逆にこれらの条件を満たすような操作後の状態は常に作ることが可能です。これらの条件を満たしているならば、1 の増加箇所と減少箇所を左から順番にマッチングさせていくことで実際に操作列を構成することができるからです。

以上の考察から、エリアで区切って各エリアに存在する 1 の個数を決めていく以下のDPを組みます。

$dp\lbrack i \rbrack\lbrack j \rbrack\lbrack k \rbrack =$ エリア $i-1$ までの 1 の個数を決め終わって、操作回数が $j$ で、これまで 1 を合計 $k$ 個使っているような状態の数

このDPの遷移は、エリア $i$ に置く 1 の個数が $one\lbrack i \rbrack$ 個未満かどうかで2通りに分けられます。それぞれの遷移で

最小操作回数で実現するパターンだけを考える。つまりあるエリア $i$ に右側から 1 を借りてきて、さらに左側にも貸すような無駄な操作を考えない。
どこから借りた/どこから貸したということは考慮せずに、個数の増減だけを考える。

という2点を守ることで、操作列が違っても結果が同じ文字列になるものは1通りとなるように数えることができます。

エリア $i$ に置く `1` の個数が $one\lbrack i \rbrack$ 個未満

このときは元々あった $one\lbrack i \rbrack$ 個の 1 のうち一部をエリア $i$ で使い、残りの 1 を左側に貸しているという扱いになります。

使う 1 の個数は複数通り考えられるので、 $dp\lbrack i \rbrack\lbrack j \rbrack\lbrack k \rbrack$ からの遷移先の候補は

$dp\lbrack i+1 \rbrack\lbrack j \rbrack\lbrack k \rbrack$
$dp\lbrack i+1 \rbrack\lbrack j \rbrack\lbrack k+1 \rbrack$
…
$dp\lbrack i+1 \rbrack\lbrack j \rbrack\lbrack k+one\lbrack i \rbrack -1 \rbrack$

となります。これを全部試していると全体で4乗になってしまうので、貰うDPにして累積和を使います。

$dp\lbrack i+1 \rbrack\lbrack j \rbrack\lbrack k \rbrack$ に遷移する値を累積和で計算できたとします。ここで先ほどの「エリア $i$ 以前（エリア $i$ を含む）に存在する 1 の個数は $sum\lbrack i \rbrack$ 個以上である」という条件に従い、 $k \ge sum\lbrack i \rbrack$ の時だけ実際に遷移すれば良いです。

エリア $i$ に置く `1` の個数が $one\lbrack i \rbrack$ 個以上

このときは、元々あった $one\lbrack i \rbrack$ 個の 1 は必ず全てエリア $i$ で使い、さらに右側から0個以上の 1 を借りてくるという扱いになります。

借りてくる 1 の個数が複数通り考えられるので、 $dp\lbrack i \rbrack\lbrack j \rbrack\lbrack k \rbrack$ からの遷移先の候補は

$dp\lbrack i+1 \rbrack\lbrack j \rbrack\lbrack k+ one\lbrack i \rbrack \rbrack$
$dp\lbrack i+1 \rbrack\lbrack j+1 \rbrack\lbrack k+ one\lbrack i \rbrack + 1 \rbrack$
…

とこちらもそのままでは4乗になります。これの対処は斜めに累積和を取っても良いですし、 $dp\lbrack i \rbrack\lbrack j \rbrack\lbrack k \rbrack$ から

$dp\lbrack i \rbrack\lbrack j+1 \rbrack\lbrack k+1 \rbrack$ に遷移（1 を1個借りてきて同じ $i$ に遷移し、後で↓の遷移をしてもらう）
$dp\lbrack i+1 \rbrack\lbrack j \rbrack\lbrack k+ one\lbrack i \rbrack \rbrack$ に遷移（元々あった $one\lbrack i \rbrack$ 個の 1 を追加して $i+1$ に遷移）