1920E - Counting Binary Strings - CodeForces Solution

#include<bits/stdc++.h>
#define bug(x) cerr<<#x<<" = "<<x<<'\n'
using namespace std;
using ll = long long;

const int mod = 998244353;
const int maxn = 2550;

int add(int a, int b) {
    return (a + b > mod ? a + b - mod : a + b);
}

int mul(int a, int b) {
    return 1ll * a * b % mod;
}

int N, K;
int dp[maxn][maxn];

void solve() {
    cin >> N >> K;
    for (int i = 1; i <= K; ++i) {
        dp[i][i - 1] = 1;
    }
    int ans = 0;
    for (int i = 1; i <= N; ++i) {
        for (int j = 0; j < K; ++j) if (dp[i][j]) {
            // if (i == 2) {
            //     cerr<<'\n';
            //     bug(j);
            //     bug(dp[i][j]);
            // }
            for (int t = 0; i + (j + 1) * t <= N; ++t) if (t <= K - 1 - j) {
                int ni = i + (j + 1) * t + (t + 1);
                if (ni <= N) 
                    dp[ni][t] = add(dp[ni][t], dp[i][j]);
                ni -= (t + 1);
                if (ni == N && t) {
                    ans = add(ans, dp[i][j]);
                    // bug(dp[i][j]);
                    // bug(t);
                    // bug(i + (j + 1) * t);
                    // bug(i);
                    // bug(j);
                    // cerr<<'\n';
                }
            }
        } 
    }
    for (int i = 0; i < K; ++i) ans = add(ans, dp[N][i]);
    cout<<ans<<'\n';
    for (int i = 0; i <= N; ++i)
        for (int j = 0; j <= K; ++j) 
             dp[i][j] = 0;
}

int main() {
    ios::sync_with_stdio(0); cin.tie(0);
    //freopen("gg.in", "r", stdin);
    //freopen("gg.out", "w", stdout); 
    int T; cin >> T; while (T--) solve();
}