Problem p02634 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using ull = unsigned long long;
#define REP(i, n) for (ll i = 0; i < (n); ++i)
#define RREP(i, n) for (ll i = (n)-1; i >= 0; --i)
#define ALL(v) (v).begin(), (v).end()
template <class T>
using reverse_priority_queue = priority_queue<T, vector<T>, greater<T>>;

struct mint {
  static const long long MOD = 998244353;

  long long v;

  mint() : v(0) {}
  mint(long long x) {
    x = x % MOD;
    if (x < 0) {
      x += MOD;
    }
    v = x;
  }
  mint &operator+=(mint a) {
    v += a.v;
    if (v >= MOD) {
      v -= MOD;
    }
    return *this;
  }
  mint &operator*=(mint a) {
    v *= a.v;
    v %= MOD;
    return *this;
  }
  mint &operator-=(mint a) {
    v -= a.v;
    if (v < 0) {
      v += MOD;
    }
    return *this;
  }
  mint &operator/=(mint a) { return (*this) *= a.inv(); }

  mint operator*(mint a) const { return mint(v) *= a; }

  mint operator+(mint a) const { return mint(v) += a; }

  mint operator-(mint a) const { return mint(v) -= a; }

  mint operator/(mint a) const { return mint(v) /= a; }

  mint pow(long long k) const {
    mint res(1), tmp(v);
    while (k) {
      if (k & 1)
        res *= tmp;
      tmp *= tmp;
      k >>= 1;
    }
    return res;
  }
  mint inv() { return pow(MOD - 2); }

  static mint comb(long long n, int k) {
    mint res(1);
    for (int i = 0; i < k; ++i) {
      res *= mint(n - i);
      res /= mint(i + 1);
    }
    return res;
  }

  static mint factorial(long long n) {
    mint res(1);
    for (int i = n; i > 1; --i) {
      res *= mint(i);
    }
    return res;
  }

  bool operator<(const mint &a) const { return v < a.v; };
};

mint operator*(long long l, mint r) { return r * l; }

std::ostream &operator<<(std::ostream &stream, const mint &m) {
  stream << m.v;
  return stream;
}

ll a, b, c, d;
// map<array<ll, 3>, mint> cache;
mint val[3001][3001][3];
bool flg[3001][3001][3];
mint solve(ll x, ll y, ll n) {
  if (x == a && y == b) {
    return (n == 0);
  } else if (y == b && (n == 0 || n == 2)) {
    return 0;
  } else if (x > a && n == 0) {
    return 0;
  }
  if (flg[x][y][n]) {
    return val[x][y][n];
  }

  mint ret;
  if (n == 0) {
    assert(x == a);
    ret = mint(x - 1).pow(y - b);
  } else if (n == 1) {
    if (x > a) {
      ret = y * (solve(x - 1, y, 0) + solve(x - 1, y, 1) + solve(x - 1, y, 2));
    } else {
      ret = solve(x, y - 1, 0) + (x - 1) * solve(x, y - 1, 1);
    }
  } else {
    assert(n == 2);
    ret = solve(x, y - 1, 1) + x * solve(x, y - 1, 2);
  }
  val[x][y][n] = ret;
  flg[x][y][n] = true;
  return ret;
}

int main() {
  cin >> a >> b >> c >> d;
  memset(flg, 0, sizeof(flg));
  cout << solve(c, d, 0) + solve(c, d, 1) + solve(c, d, 2) << endl;
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using ull = unsigned long long;
#define REP(i, n) for (ll i = 0; i < (n); ++i)
#define RREP(i, n) for (ll i = (n)-1; i >= 0; --i)
#define ALL(v) (v).begin(), (v).end()
template <class T>
using reverse_priority_queue = priority_queue<T, vector<T>, greater<T>>;

struct mint {
  static const long long MOD = 998244353;

  long long v;

  mint() : v(0) {}
  mint(long long x) {
    x = x % MOD;
    if (x < 0) {
      x += MOD;
    }
    v = x;
  }
  mint &operator+=(mint a) {
    v += a.v;
    if (v >= MOD) {
      v -= MOD;
    }
    return *this;
  }
  mint &operator*=(mint a) {
    v *= a.v;
    v %= MOD;
    return *this;
  }
  mint &operator-=(mint a) {
    v -= a.v;
    if (v < 0) {
      v += MOD;
    }
    return *this;
  }
  mint &operator/=(mint a) { return (*this) *= a.inv(); }

  mint operator*(mint a) const { return mint(v) *= a; }

  mint operator+(mint a) const { return mint(v) += a; }

  mint operator-(mint a) const { return mint(v) -= a; }

  mint operator/(mint a) const { return mint(v) /= a; }

  mint pow(long long k) const {
    mint res(1), tmp(v);
    while (k) {
      if (k & 1)
        res *= tmp;
      tmp *= tmp;
      k >>= 1;
    }
    return res;
  }
  mint inv() { return pow(MOD - 2); }

  static mint comb(long long n, int k) {
    mint res(1);
    for (int i = 0; i < k; ++i) {
      res *= mint(n - i);
      res /= mint(i + 1);
    }
    return res;
  }

  static mint factorial(long long n) {
    mint res(1);
    for (int i = n; i > 1; --i) {
      res *= mint(i);
    }
    return res;
  }

  bool operator<(const mint &a) const { return v < a.v; };
};

mint operator*(long long l, mint r) { return r * l; }

std::ostream &operator<<(std::ostream &stream, const mint &m) {
  stream << m.v;
  return stream;
}

ll a, b, c, d;
// map<array<ll, 3>, mint> cache;
unordered_map<ll, mint> cache;
mint solve(ll x, ll y, ll n) {
  if (x == a && y == b) {
    return (n == 0);
  } else if (y == b && (n == 0 || n == 2)) {
    return 0;
  } else if (x > a && n == 0) {
    return 0;
  }
  const ll key = n + 3 * y + 3 * 3001 * x;
  // if(cache.count({x, y, n})){
  //     return cache.at({x, y, n});
  // }
  if (cache.count(key)) {
    return cache.at(key);
  }

  mint ret;
  if (n == 0) {
    assert(x == a);
    ret = mint(x - 1).pow(y - b);
  } else if (n == 1) {
    if (x > a) {
      ret = y * (solve(x - 1, y, 0) + solve(x - 1, y, 1) + solve(x - 1, y, 2));
    } else {
      ret = solve(x, y - 1, 0) + (x - 1) * solve(x, y - 1, 1);
    }
  } else {
    assert(n == 2);
    ret = solve(x, y - 1, 1) + x * solve(x, y - 1, 2);
  }
  // cache[{x, y, n}] = ret;
  cache[key] = ret;
  return ret;
}

int main() {
  cin >> a >> b >> c >> d;
  cout << solve(c, d, 0) + solve(c, d, 1) + solve(c, d, 2) << endl;
  return 0;
}