ECR157-F Fancy Arrays - zawatinはてなブログ

問題

https://codeforces.com/contest/1895/problem/F (diff: 2600)

$n, x, k$ が与えられる。以下の条件を満たす長さ $n$ の非負整数列 $a$ を数え上げよ。

$x, x + 1, x + 2, \dots, x + k - 1$ のうち少なくとも一つの要素が少なくとも一回登場する
$i = 2, 3, \dots, n$ について $|a_{i} - a_{i - 1}|\ \le\ k$

制約

$1\ \le\ n, k\ \le\ 10^{9}$
$0\ \le\ x\ \le\ 40$

解法

$a$ は次の条件を全て満たすとき、またそのときに限り、「二つ目の条件を満たすけど、一つ目の条件を満たさない数列」となる

二つ目の条件を満たす
「全ての要素が $0$ 以上 $x - 1$ 以下である」または「全ての要素が $x + k$ 以上」

このことから、 $f(l, r)$ を $l$ 以上 $r$ 以下の整数からなる二つ目の条件を満たす数列の数、 $M$ を十分大きい数として元の問題の解が $f(0, M) - f(0, x - 1) - f(x + k, M)$ であることがわかった。

$f(l, r)$ を求める方法を考える。二つの解法を挙げる。

一つ目の方法: dp

$\text{dp}_{i, j}$ を長さ $i$ の、末尾の項が $j$ であるような数列の数とすると、 $O((r - l)n)$ でこの問題を解くことができる。このdpの計算に行列累乗を使うと、 $O((r - l)^{3}\log n)$ で解くこともできる。

二つ目の方法: 差分の決め方を全探索

$a_{i + 1} - a_{i} = d_{i}$ とすると、 $d$ の数は $(2k + 1)^{n - 1}$ 通りある。この数列全てに対してそれぞれ適切な $a_{1}$ の数を数えれば良い。

$f(0, x - 1)$ は「一つ目の方法」の行列累乗を用いて計算できる。

$f(0, M)$ 及び $f(x + k, M)$ の計算が難しい。そこで「二つ目の方法」の適切な $a_{1}$ についてちゃんと考えてみる。

$d$ を固定すると、 $f(0, M)$ 、 $f(x + k, M)$ について $a_{1}$ としてありうる最大値が等しい。

$d$ を固定すると、 $a_{1}$ としてありうる最小値は $f(0, M)$ より $f(x + k, M)$ の方が丁度 $x + k$ だけ大きいことが分かる。

これが全ての $d$ について言える。

このことから、 $f(0, M)$ と $f(x + k, M)$ の計算は依然できていないが、 $f(0, M) - f(x + M) = (x + k)(2k + 1)^{n - 1}$ であることがわかった。

以上より $f(0, M) - f(0, x - 1) - f(x + k, M)$ を計算することに成功した。

実装

#include <bits/stdc++.h>

#include "Src/Template/IOSetting.hpp"
using namespace zawa;

#include "atcoder/modint"
using mint = atcoder::modint1000000007;

using Mat = std::vector<std::vector<mint>>;
Mat mult(const Mat& a, const Mat& b) {
    Mat res(a.size(), std::vector<mint>(b[0].size()));
    for (int i{} ; i < (int)a.size() ; i++) {
        for (int j{} ; j < (int)b[0].size() ; j++) {
            for (int k{} ; k < (int)a[0].size() ; k++) {
                res[i][j] += a[i][k] * b[k][j];
            }
        }
    }
    return res;
}

Mat pow(const Mat& a, int k) {
    Mat res(a.size(), std::vector<mint>(a[0].size()));
    for (int i{} ; i < (int)a[0].size() ; i++) res[i][i] = mint::raw(1);
    Mat base{a};
    while (k) {
        if (k & 1) {
            res = mult(res, base);
        }
        base = mult(base, base);
        k >>= 1;
    }
    return res;
}

mint solve() {
    int n, x, k; std::cin >> n >> x >> k;
    mint ans{};
    if (x > 0) {
        Mat A(x, std::vector<mint>(x));
        for (int i{} ; i < x ; i++) {
            for (int j{} ; j < x ; j++) {
                if (std::abs(i - j) <= k) {
                    A[i][j] = mint::raw(1);
                }
            }
        }
        A = pow(A, n - 1);
        for (int i{} ; i < x ; i++) {
            for (int j{} ; j < x ; j++) {
                ans += A[i][j];
        //         std::cout << A[i][j].val() << (j + 1 == x ? '\n' : ' ');
            }
        }
    }
    ans = mint{x + k} * mint{2 * k + 1}.pow(n - 1) - ans;
    return ans;
}

int main() {
    SetFastIO();

    int t; std::cin >> t;
    while (t--) {
        std::cout << solve().val() << '\n';
    }
}

感想

バチャ中は $x, \dots, x + k - 1$ が全て少なくとも一回登場する列を数える問題と勘違いしていた。でもこの勘違いが無くても解けなかっただろう。

最初の言い換えから見つけるのが難しい。どうやったら気付けるようになれるのかなぁ。

毎回行列累乗をソラ書きしているのアホらしいのでやめたい。盆栽書かないと。