间距数列

题目

满足以下条件的长为 \(n\) 的整数序列 \(a_1, a_2, \cdots, a_n\) 有多少个？

• \(1 \leqslant a_i \leqslant m\)，\(1 \leqslant i \leqslant n\)
• \(|a_i-a_{i+1}| \geqslant k\ (1 \leqslant i \leqslant n-1)\)

输出答案除以 \(998244353\) 的余数。

限制：

• \(2 \leqslant n \leqslant 1000\)
• \(1 \leqslant m \leqslant 5000\)
• \(0 \leqslant k \leqslant m-1\)

算法分析

记 dp[i][j] 表示长为 \(i\) 的结尾等于 \(j\) 的数列个数

转移方程：

\( dp[i][j] = (dp[i-1][1] + \cdots + dp[i-1][j-k]) + (dp[i-1][j+k] + \cdots + dp[i-1][m]) \)

然后用前缀和去优化
记 s[i][j] 表示长为 \(i\) 且结尾 \(\leqslant j\) 的数列个数

初始值：

\(dp[1][1 \sim m] = 1\)
\(s[1][j] = j\)

还需特判一下 \(k = 0\) 的情况，答案是 \(m^n\)

代码实现

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 1; i <= (n); ++i)

using namespace std;
using ll = long long;

const int mod = 998244353;
//const int mod = 1000000007;
struct mint {
    ll x;
    mint(ll x=0):x((x%mod+mod)%mod) {}
    mint operator-() const {
        return mint(-x);
    }
    mint& operator+=(const mint a) {
        if ((x += a.x) >= mod) x -= mod;
        return *this;
    }
    mint& operator-=(const mint a) {
        if ((x += mod-a.x) >= mod) x -= mod;
        return *this;
    }
    mint& operator*=(const mint a) {
        (x *= a.x) %= mod;
        return *this;
    }
    mint operator+(const mint a) const {
        return mint(*this) += a;
    }
    mint operator-(const mint a) const {
        return mint(*this) -= a;
    }
    mint operator*(const mint a) const {
        return mint(*this) *= a;
    }
    mint pow(ll t) const {
        if (!t) return 1;
        mint a = pow(t>>1);
        a *= a;
        if (t&1) a *= *this;
        return a;
    }

    // for prime mod
    mint inv() const {
        return pow(mod-2);
    }
    mint& operator/=(const mint a) {
        return *this *= a.inv();
    }
    mint operator/(const mint a) const {
        return mint(*this) /= a;
    }
};
istream& operator>>(istream& is, mint& a) {
    return is >> a.x;
}
ostream& operator<<(ostream& os, const mint& a) {
    return os << a.x;
}

mint dp[1005][5005];
mint s[1005][5005];

int main() {
    int n, m, k;
    cin >> n >> m >> k;
    
    if (k == 0) {
        cout << mint(m).pow(n) << '\n';
        return 0;
    }
    
    rep(j, m) dp[1][j] = 1;
    rep(j, m) s[1][j] = j;
    for (int i = 2; i <= n; ++i) {
        for (int j = 1; j <= m; ++j) {
            mint now;
            if (j-k >= 0) now += s[i-1][j-k];
            if (j+k-1 <= m) now += s[i-1][m] - s[i-1][j+k-1];
            dp[i][j] = now;
            s[i][j] = s[i][j-1] + now;
        }
    }
    
    mint ans;
    rep(j, m) {
        ans += dp[n][j];
    }
    
    cout << ans << '\n';
    
    return 0;
}