1051D - Bicolorings - CodeForces Solution

#include<bits/stdc++.h>
#define ll long long
#define fi first
#define se second
#define IOS ios::sync_with_stdio(false); cin.tie(0); cout.tie(0);
#define all(x) x.begin(), x.end()
#define pb push_back
#define pf push_front
#define bit(n, x) ((n >> x) & 1)

using namespace std;
double begintime, endtime;
inline void CALC_TIME()
{
    #ifndef ONLINE_JUDGE
    endtime = clock();
    cerr << "\nexecution time : " << (endtime - begintime + 1) / 1000 << " s\n";
    #endif
}
inline void file(string NAME)
{
    #ifndef ONLINE_JUDGE
    freopen((NAME + ".INP").c_str(), "r", stdin);
    freopen((NAME + ".OUT").c_str(), "w", stdout);
    #endif
}
const int MOD = 998244353;
const int maxN = 1e6 + 68;
const int INF = 1e9;
int n, k;
ll dp[1005][1 << 3][2005];
ll dq(int i, int mask, int comp)
{
    if (i == n - 1) return dp[i][mask][comp] = (comp == k);
    if (dp[i][mask][comp] != -1) return dp[i][mask][comp];
    ll res = 0;
    for (int next_mask = 0; next_mask < (1 << 2); next_mask++) {
        if ((mask ^ next_mask) == 0) res += dq(i + 1, next_mask, comp);
        else if ((mask ^ next_mask) == 3) {
            if (mask == 0 && next_mask == 3) res += dq(i + 1, next_mask, comp + 1);
            else if (mask == 3 && next_mask == 0) res += dq(i + 1, next_mask, comp + 1);
            else if (mask == 1 && next_mask == 2) res += dq(i + 1, next_mask, comp + 2);
            else if (mask == 2 && next_mask == 1) res += dq(i + 1, next_mask, comp + 2);
        }
        else if ((mask ^ next_mask) == 1) {
            if (mask == 0 && next_mask == 1) res += dq(i + 1, next_mask, comp + 1);
            else if (mask == 1 && next_mask == 0) res += dq(i + 1, next_mask, comp);
            else if (mask == 2 && next_mask == 3) res += dq(i + 1, next_mask, comp);
            else if (mask == 3 && next_mask == 2) res += dq(i + 1, next_mask, comp + 1);
        }
        else if ((mask ^ next_mask) == 2) {
            if (mask == 0 && next_mask == 2) res += dq(i + 1, next_mask, comp + 1);
            else if (mask == 2 && next_mask == 0) res += dq(i + 1, next_mask, comp);
            else if (mask == 1 && next_mask == 3) res += dq(i + 1, next_mask, comp);
            else if (mask == 3 && next_mask == 1) res += dq(i + 1, next_mask, comp + 1);
        }
        res %= MOD;
    }
    res %= MOD;
    return dp[i][mask][comp] = res;
}
string get(int x)
{
    string res = "";
    for (int i = 0; i < 2; ++i) {
        if (bit(x, i)) res += "1";
        else res += "0";
    }
    return res;
}
signed main() {
    begintime = clock();
    IOS
    cin >> n >> k;
    memset(dp, -1, sizeof dp);
    ll ans = 0;
    for (int mask = 0; mask < (1 << 2); mask++) {
        if (bit(mask, 0) != bit(mask, 1)) ans += dq(0 , mask, 2);
        else ans += dq(0 , mask, 1);
        ans %= MOD;
    }
    cout << ans;
    return 0;
}