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

#include <iostream>

using namespace std;

typedef long long ll;

const int MOD = 1e9 + 7; // change if needed

struct Modint {
  ll val;

  Modint(ll _val = 0) : val(_val % MOD) {}

  Modint operator+(Modint other) const { return Modint(val + other.val); }

  void operator+=(Modint other) {
    val += other.val;
    val %= MOD;
  }

  Modint operator-() const { return Modint(MOD - val); }

  Modint operator-(Modint other) const { return Modint(val + MOD - other.val); }

  void operator-=(Modint other) {
    val += MOD - other.val;
    val %= MOD;
  }

  Modint operator*(Modint other) const { return Modint(val * other.val); }

  void operator*=(Modint other) {
    val *= other.val;
    val %= MOD;
  }

  bool operator==(Modint other) const { return val == other.val; }

  bool operator!=(Modint other) const { return val != other.val; }
};

Modint exp(Modint a, int k) {
  if (k == 0) {
    return Modint(1);
  } else if (k % 2 == 0) {
    Modint half = exp(a, k / 2);
    return half * half;
  } else {
    return a * exp(a, k - 1);
  }
}

Modint inv(Modint a) { return exp(a, MOD - 2); }

ostream &operator<<(ostream &out, Modint p) {
  out << p.val;
  return out;
}

const int MAX_N = 5e3 + 5;

Modint dp[MAX_N][MAX_N];
Modint isum[MAX_N];
Modint dsum[MAX_N][MAX_N];

Modint irange(int l, int r) {
  l = max(0, l);
  if (r < l)
    return Modint(0);
  if (l == 0)
    return isum[r];
  return isum[r] - isum[l - 1];
}

Modint dget(int diag, int idx) {
  if (diag < 0 || idx < 0)
    return Modint(0);
  return dsum[diag][idx];
}

Modint tpow[MAX_N];

int main() {
  tpow[0] = Modint(1);
  for (int i = 1; i < MAX_N; i++) {
    tpow[i] = Modint(2) * tpow[i - 1];
  }

  int n, A, B;
  cin >> n >> A >> B;

  if (A > B)
    swap(A, B);
  if (A == 1) {
    cout << exp(Modint(2), n) << endl;
    return 0;
  }

  dp[0][0] = Modint(1);
  dsum[0][0] = Modint(1);
  isum[0] = Modint(1);
  for (int i = 1; i <= n + 1; i++) {
    for (int b = 1; b <= i; b++) {
      if (b == 1) {
        if (i == 1) {
          dp[i][b] = 1;
        } else {
          dp[i][b] = irange(i - A, i - 2);
        }
      } else {
        dp[i][b] = dget(i - b, i - A - 1);
        if (b - 1 < B) {
          dp[i][b] += dp[i - 1][b - 1];
        }
      }

      if (b < B) {
        isum[i] += dp[i][b];
        dsum[i - b][i] += dp[i][b];
      }
      dsum[i - b][i] += dsum[i - b][i - 1];
    }

    isum[i] += isum[i - 1];
  }

  /*
  for (int i = 1; i <= n + 1; i++) {
    for (int b = 1; b <= n + 1; b++) {
      cout << dp[i][b] << " ";
    }
    cout << endl;
  }
  */

  Modint ans(0);
  for (int i = 1; i <= n + 1; i++) {
    for (int b = B; b <= n + 1; b++) {
      if (i == n + 1) {
        if (b != B) {
          ans += dp[i][b];
        }
      } else {
        ans += dp[i][b] * tpow[n - i];
      }
    }
  }
  cout << ans << endl;
}


Error: TLE

Fixed submission:

#include <iostream>

using namespace std;

typedef long long ll;

const int MOD = 1e9 + 7; // change if needed

struct Modint {
  ll val;

  Modint(ll _val = 0) : val(_val % MOD) {}

  Modint operator+(Modint other) const { return Modint(val + other.val); }

  void operator+=(Modint other) {
    val += other.val;
    val %= MOD;
  }

  Modint operator-() const { return Modint(MOD - val); }

  Modint operator-(Modint other) const { return Modint(val + MOD - other.val); }

  void operator-=(Modint other) {
    val += MOD - other.val;
    val %= MOD;
  }

  Modint operator*(Modint other) const { return Modint(val * other.val); }

  void operator*=(Modint other) {
    val *= other.val;
    val %= MOD;
  }

  bool operator==(Modint other) const { return val == other.val; }

  bool operator!=(Modint other) const { return val != other.val; }
};

Modint exp(Modint a, int k) {
  if (k == 0) {
    return Modint(1);
  } else if (k % 2 == 0) {
    Modint half = exp(a, k / 2);
    return half * half;
  } else {
    return a * exp(a, k - 1);
  }
}

Modint inv(Modint a) { return exp(a, MOD - 2); }

ostream &operator<<(ostream &out, Modint p) {
  out << p.val;
  return out;
}

const int MAX_N = 5e3 + 5;

Modint dp[MAX_N][MAX_N];
Modint isum[MAX_N];
Modint dsum[MAX_N][MAX_N];

Modint irange(int l, int r) {
  l = max(0, l);
  if (r < l)
    return Modint(0);
  if (l == 0)
    return isum[r];
  return isum[r] - isum[l - 1];
}

Modint dget(int diag, int idx) {
  if (diag < 0 || idx < 0)
    return Modint(0);
  return dsum[diag][idx];
}

int main() {
  int n, A, B;
  cin >> n >> A >> B;

  if (A > B)
    swap(A, B);
  if (A == 1) {
    cout << exp(Modint(2), n) << endl;
    return 0;
  }

  dp[0][0] = Modint(1);
  dsum[0][0] = Modint(1);
  isum[0] = Modint(1);
  for (int i = 1; i <= n + 1; i++) {
    for (int b = 1; b <= i; b++) {
      if (b == 1) {
        if (i == 1) {
          dp[i][b] = 1;
        } else {
          dp[i][b] = irange(i - A, i - 2);
        }
      } else {
        dp[i][b] = dget(i - b, i - A - 1);
        if (b - 1 < B) {
          dp[i][b] += dp[i - 1][b - 1];
        }
      }

      if (b < B) {
        isum[i] += dp[i][b];
        dsum[i - b][i] += dp[i][b];
      }
      dsum[i - b][i] += dsum[i - b][i - 1];
    }

    isum[i] += isum[i - 1];
  }

  /*
  for (int i = 1; i <= n + 1; i++) {
    for (int b = 1; b <= n + 1; b++) {
      cout << dp[i][b] << " ";
    }
    cout << endl;
  }
  */

  Modint ans(0);
  for (int i = 1; i <= n + 1; i++) {
    for (int b = B; b <= n + 1; b++) {
      if (i == n + 1) {
        if (b != B) {
          ans += dp[i][b];
        }
      } else {
        ans += dp[i][b] * exp(Modint(2), n - i);
      }
    }
  }
  cout << ans << endl;
}