Codeforces Round #554 F. Neko Rules the Catniverse

問題概要

整数 $n, k, m$ が与えられる。以下の条件を満たす要素数 $k$ の数列 $a_{1}, ..., a_{k}$ は何通りあるか、 $\bmod 10^{9}+7$ で求めよ。

$1 \le a_{i} \le n$
$a_{i} \le a_{i-1} + m$
$a_{i}$ は相異なる

制約

F1とF2での範囲が異なる。
- F1: $1 \le n \le 10^{5}$
- F2: $1 \le n \le 10^{9}$
$1 \le k \le \min(n, 12)$
$1 \le m \le 4$

解法

何となく数列の頭から順に決めていくDPを考えたくなりますが、「相異なる」という条件が厄介です。これをまともに扱おうとそれまで使ったことのある整数を情報として持つことが必要になり、数列の長さがたった12しかないとはいえ値の種類数 $n$ がそこそこ大きいのでとても大変です。

見方を変えて、「値のほうを $1$ から $n$ まで順に見ていって、使う整数を数列に挿入していく」というDPを考えてみましょう。以下のようなDPテーブルを考えます。

$dp\lbrack i \rbrack\lbrack j\rbrack =$ 整数 $1$ から $i$ までを使うかどうか決めて、長さが $j$ で、問題の条件を満たしているような数列の個数

実はこれだけでは情報として不十分ですが、いったんこんな方針だと思ってみましょう。

$dp\lbrack i \rbrack\lbrack j\rbrack$ からの遷移を考えます。 $i+1$ を挿入するとき、それを挿入できる場所は

先頭
既に数列にある整数のうち、 $i+1 \le x + m$ であるような $x$ の後ろ

のいずれかです。

また、挿入する位置の後ろの整数との関係は常に条件を満たします。後ろの整数を $y$ とすると条件式は $y \le i+1+m$ ですが、今までの値より大きい値を常に挿入するのでこれは常に満たされます。

f:id:betrue12:20190425220354p:plain

そして数列が具体的にどのような並びであろうと、挿入できる箇所の個数は変わりません。ここが挿入DPの大きな特徴です。

そのため $i+1$ を挿入できる場所の個数を知るためには、 $i+1-m$ から $i$ までの $m$ 個の整数が数列に存在しているかという情報があれば十分です。これをDPテーブルの状態に加えましょう。

$dp\lbrack i \rbrack\lbrack j\rbrack\lbrack S \rbrack =$ 整数 $1$ から $i$ までを使うかどうか決めて、長さが $j$ で、 $i+1-m$ から $i$ までの整数のうち数列に含まれているものの集合が $S$ であって、問題の条件を満たしているような数列の個数

ここで $S$ は $m$ 個の整数のある/なしを表しているので $2^{m}$ 通りです。実装上はそれぞれをビットで管理した整数として扱いましょう。最下位ビットを整数 $i$ に、最上位ビットを整数 $i+1-m$ に対応させておくと、ビットシフトで見ている範囲をずらすことができます。

こう置くと、 $dp\lbrack i \rbrack\lbrack j\rbrack\lbrack S \rbrack$ からの遷移は

整数を使う。
- $S$ の立っているビットの個数を $b$ として、 $b+1$ 倍して $dp\lbrack i+1 \rbrack\lbrack j+1\rbrack\lbrack (S\lt\lt 1) \% 2^{m} + 1 \rbrack$ に遷移する。
整数を使わない。
- $dp\lbrack i+1 \rbrack\lbrack j\rbrack\lbrack (S\lt\lt 1) \% 2^{m}\rbrack$ に遷移する。

と書くことができます。

※一般に挿入DPで考慮すべきこととして、「隣り合った要素がいったん条件に違反していても、後で他の値を間に入れることで違反が解消されるかもしれない」というものがあります。しかし今回は常に今までより大きな値を入れていくため、 $a_{i} \le a_{i-1} + m$ に違反しているような $a_{i-1}, a_{i}$ の間にこれらより大きな値を入れても依然違反状態であり、解消されることはないので考慮しなくて大丈夫です。

このDPを普通に計算すると計算量が $O(nkm2^{m})$ になります。F1の制約においてはこれは十分間に合います。

F2においては $n$ が非常に大きく、 $O(n)$ が掛かるような解法は厳しいです。これを解くためDPの遷移に注目すると、

$dp\lbrack i+1 \rbrack\lbrack j\rbrack\lbrack S \rbrack$ は、いくつかの $dp\lbrack i \rbrack\lbrack \cdot \rbrack\lbrack \cdot \rbrack$ に $i$ に依存しない係数を掛けたものの和（線形結合）である。
$dp\lbrack i \rbrack\lbrack \cdot \rbrack\lbrack \cdot \rbrack$ の要素数を $L$ と置くと、 $L = (k+1)2^{m}$ であり最大でも200程度と非常に小さい