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

#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <map>
#include <set>
#include <string>
#include <vector>

#include <cmath>
#include <queue>
#include <stack>

using namespace std;

#define LL long long

const int dx[4] = {1, -1, 0, 0};
const int dy[4] = {0, 0, 1, -1};

LL mod(LL x, LL md) { return ((x % md) + md) % md; }

class Mint {
public:
  static LL md;
  LL x;
  Mint(LL _x = 0) { x = mod(_x, md); }
  Mint(const Mint &y) { x = mod(y.x, md); }
};

void MintMyExgcd(LL a, LL b, LL &x, LL &y) {
  // solve equation ax + by = 1 (mod m)
  LL md = Mint::md;
  if (b == 0) {
    y = 0;
    x = 1;
    return;
  }
  MintMyExgcd(b, a % b, x, y);
  LL t = y;
  y = mod(x - (a / b) * y, md);
  x = t;
}

LL MintNY(LL x) {
  if (x == 0) {
    return 1;
  }
  LL md = Mint::md;
  LL a, b;
  MintMyExgcd(x, md, a, b);
  while (a < 0) {
    a += md;
  }
  return a;
}

/* LL gcd(LL a, LL b) {
        if (b == 0) {
                return a;
        }
        return gcd(b, a % b);
} */

Mint operator+(Mint a, Mint b) { return Mint(mod(a.x + b.x, Mint::md)); }
Mint operator-(Mint a, Mint b) { return Mint(mod(a.x - b.x, Mint::md)); }
Mint operator*(Mint a, Mint b) { return Mint(mod(a.x * b.x, Mint::md)); }
Mint operator/(Mint a, Mint b) {
  Mint c = Mint(MintNY(b.x));
  return Mint(mod(a.x * c.x, Mint::md));
}

Mint operator-(Mint a) { return Mint(-a.x); }
bool operator==(Mint a, Mint b) { return a.x == b.x; }
bool operator!=(Mint a, Mint b) { return a.x != b.x; }
bool operator>(Mint a, Mint b) { return a.x > b.x; }
bool operator<(Mint a, Mint b) { return a.x < b.x; }
bool operator>=(Mint a, Mint b) { return a.x >= b.x; }
bool operator<=(Mint a, Mint b) { return a.x <= b.x; }

void operator+=(Mint &a, Mint b) { a = a + b; }
void operator-=(Mint &a, Mint b) { a = a - b; }
void operator*=(Mint &a, Mint b) { a = a * b; }
void operator/=(Mint &a, Mint b) { a = a / b; }

ostream &operator<<(ostream &out, Mint a) {
  out << a.x;
  return out;
}

istream &operator>>(istream &in, Mint &a) {
  in >> a.x;
  return in;
}

Mint factorial(LL n) {
  if (n == 0) {
    return Mint(1);
  }
  Mint res = 1;
  for (LL i = 1; i <= n; ++i) {
    res *= i;
  }
  return res;
}
Mint multiRange(LL s, LL t) {
  Mint res = 1;
  for (LL i = s; i <= t; ++i) {
    res *= i;
  }
  return res;
}

Mint getC(LL n, LL k) {
  if (n - k < k) {
    return getC(n, n - k);
  }
  return multiRange(n - k + 1, n) / factorial(k);
}

Mint getA(LL n, LL k) { return multiRange(n - k + 1, n); }

Mint countNonNegativeDivision(LL n, LL k) { return getC(n + k - 1, k - 1); }
Mint countPositiveDivision(LL n, LL k) {
  if (n < k) {
    return 0;
  }
  if (k <= 0) {
    return 0;
  }
  return getC(n - 1, k - 1);
}
Mint quickMod(Mint a, LL b) {
  Mint res = 1;
  while (b > 0) {
    if (b % 2) {
      res *= a;
    }
    a *= a;
    b /= 2;
  }
  return res;
}
template <class TypeName> vector<TypeName> getVector1D(int n, TypeName val) {
  return vector<TypeName>(n, val);
}
template <class TypeName>
vector<vector<TypeName>> getVector2D(int n1, int n2, TypeName val) {
  vector<vector<TypeName>> res(n1);
  for (int i = 0; i < n1; ++i) {
    res[i] = vector<TypeName>(n2, val);
  }
  return res;
}
template <class TypeName>
vector<vector<vector<TypeName>>> getVector3D(int n1, int n2, int n3,
                                             TypeName val) {
  vector<vector<vector<TypeName>>> res(n1);
  for (int i = 0; i < n1; ++i) {
    res[i] = vector<vector<TypeName>>(n2);
    for (int j = 0; j < n2; ++j) {
      res[i][j] = vector<TypeName>(n3, val);
    }
  }
  return res;
}
template <class TypeName>
vector<vector<vector<vector<TypeName>>>> getVector4D(int n1, int n2, int n3,
                                                     int n4, TypeName val) {
  vector<vector<vector<vector<TypeName>>>> res(n1);
  for (int i = 0; i < n1; ++i) {
    res[i] = vector<vector<vector<TypeName>>>(n2);
    for (int j = 0; j < n2; ++j) {
      res[i][j] = vector<vector<TypeName>>(n3);
      for (int k = 0; k < n3; ++k) {
        res[i][j][k] = vector<TypeName>(n4, val);
      }
    }
  }
  return res;
}

LL Mint::md = 1000000007;
LL md = 1000000007;

class Mat {
public:
  LL n;
  vector<vector<LL>> a;
  Mat(LL _n, bool makeIdentity = false) : n(_n) {
    a = getVector2D(n, n, (LL)(0));
    if (makeIdentity) {
      for (int i = 0; i < n; ++i) {
        a[i][i] = 1;
      }
    }
  }
  Mat(const Mat &mat) {
    n = mat.n;
    a = mat.a;
  }
  void square() {
    auto tmp = getVector2D(n, n, (LL)(0));
    for (int i = 0; i < n; ++i) {
      for (int j = 0; j < n; ++j) {
        for (int k = 0; k < n; ++k) {
          tmp[i][j] += a[i][k] * a[k][j];
          tmp[i][j] %= md;
        }
      }
    }
    // for (int i = 0; i < n; ++i) {
    //     for (int j = 0; j < n; ++j) {
    //         tmp[i][j] = mod(tmp[i][j], md);
    //     }
    // }
    swap(a, tmp);
    tmp.clear();
  }
  void multi(Mat &b) {
    auto tmp = getVector2D(n, n, (LL)(0));
    for (int i = 0; i < n; ++i) {
      for (int j = 0; j < n; ++j) {
        for (int k = 0; k < n; ++k) {
          tmp[i][j] += a[i][k] * b.a[k][j];
          tmp[i][j] %= md;
        }
      }
    }
    // for (int i = 0; i < n; ++i) {
    //     for (int j = 0; j < n; ++j) {
    //         tmp[i][j] = mod(tmp[i][j], md);
    //     }
    // }
    swap(a, tmp);
    tmp.clear();
  }
  vector<LL> vectorMulti(vector<LL> &x) {
    vector<LL> res(n, 0);
    for (int i = 0; i < n; ++i) {
      for (int j = 0; j < n; ++j) {
        res[i] += x[j] * a[j][i];
        res[i] %= md;
      }
      // res[i] = mod(res[i], md);
    }
    return res;
  }

  ~Mat() { a.clear(); }
};

Mat MatToOrderN(Mat &a, LL n) {
  Mat aa = a;
  Mat res(a.n, true);
  while (n > 0) {
    if (n & 1) {
      res.multi(aa);
    }
    aa.square();
    n /= 2;
  }
  return res;
}

LL work() {
  LL n;
  cin >> n;
  if (n < 5) {
    return 0;
  }

  Mat a(32, false);
  for (int jNew = 0; jNew <= 15; ++jNew) {
    for (int pNew = 0; pNew < 2; ++pNew) {
      int y = jNew * 2 + pNew;
      int jOld = jNew, pOld = pNew ^ 1;
      int x = jOld * 2 + pOld;
      a.a[x][y] = 1;

      if (pNew == 0 && jNew > 0) {
        for (jOld = 0; jOld < jNew; ++jOld) {
          x = jOld * 2 + 1;
          a.a[x][y] = 1;
          if (jOld > 4) {
            x = jOld * 2;
            a.a[x][y] = 1;
          }
        }
      }
    }
  }
  // for (int i = 0; i < 32; ++i) {
  //     for (int j = 0; j < 32; ++j) {
  //         cout << a.a[i][j] << ' ';
  //     }
  //     cout << endl;
  // }
  a = MatToOrderN(a, n - 5);
  vector<LL> res(32, 0);
  for (int i = 0; i <= 15; ++i) {
    res[i * 2] = 1;
  }
  res = a.vectorMulti(res);
  LL ans = mod(res[30] + res[31], md);
  return ans;
}

int main() {
  int T;
  cin >> T;
  for (int i = 0; i < T; ++i) {
    cout << work() << endl;
  }
  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <map>
#include <set>
#include <string>
#include <vector>

#include <cmath>
#include <queue>
#include <stack>

using namespace std;

#define LL long long

const int dx[4] = {1, -1, 0, 0};
const int dy[4] = {0, 0, 1, -1};

LL mod(LL x, LL md) { return ((x % md) + md) % md; }

class Mint {
public:
  static LL md;
  LL x;
  Mint(LL _x = 0) { x = mod(_x, md); }
  Mint(const Mint &y) { x = mod(y.x, md); }
};

void MintMyExgcd(LL a, LL b, LL &x, LL &y) {
  // solve equation ax + by = 1 (mod m)
  LL md = Mint::md;
  if (b == 0) {
    y = 0;
    x = 1;
    return;
  }
  MintMyExgcd(b, a % b, x, y);
  LL t = y;
  y = mod(x - (a / b) * y, md);
  x = t;
}

LL MintNY(LL x) {
  if (x == 0) {
    return 1;
  }
  LL md = Mint::md;
  LL a, b;
  MintMyExgcd(x, md, a, b);
  while (a < 0) {
    a += md;
  }
  return a;
}

/* LL gcd(LL a, LL b) {
        if (b == 0) {
                return a;
        }
        return gcd(b, a % b);
} */

Mint operator+(Mint a, Mint b) { return Mint(mod(a.x + b.x, Mint::md)); }
Mint operator-(Mint a, Mint b) { return Mint(mod(a.x - b.x, Mint::md)); }
Mint operator*(Mint a, Mint b) { return Mint(mod(a.x * b.x, Mint::md)); }
Mint operator/(Mint a, Mint b) {
  Mint c = Mint(MintNY(b.x));
  return Mint(mod(a.x * c.x, Mint::md));
}

Mint operator-(Mint a) { return Mint(-a.x); }
bool operator==(Mint a, Mint b) { return a.x == b.x; }
bool operator!=(Mint a, Mint b) { return a.x != b.x; }
bool operator>(Mint a, Mint b) { return a.x > b.x; }
bool operator<(Mint a, Mint b) { return a.x < b.x; }
bool operator>=(Mint a, Mint b) { return a.x >= b.x; }
bool operator<=(Mint a, Mint b) { return a.x <= b.x; }

void operator+=(Mint &a, Mint b) { a = a + b; }
void operator-=(Mint &a, Mint b) { a = a - b; }
void operator*=(Mint &a, Mint b) { a = a * b; }
void operator/=(Mint &a, Mint b) { a = a / b; }

ostream &operator<<(ostream &out, Mint a) {
  out << a.x;
  return out;
}

istream &operator>>(istream &in, Mint &a) {
  in >> a.x;
  return in;
}

Mint factorial(LL n) {
  if (n == 0) {
    return Mint(1);
  }
  Mint res = 1;
  for (LL i = 1; i <= n; ++i) {
    res *= i;
  }
  return res;
}
Mint multiRange(LL s, LL t) {
  Mint res = 1;
  for (LL i = s; i <= t; ++i) {
    res *= i;
  }
  return res;
}

Mint getC(LL n, LL k) {
  if (n - k < k) {
    return getC(n, n - k);
  }
  return multiRange(n - k + 1, n) / factorial(k);
}

Mint getA(LL n, LL k) { return multiRange(n - k + 1, n); }

Mint countNonNegativeDivision(LL n, LL k) { return getC(n + k - 1, k - 1); }
Mint countPositiveDivision(LL n, LL k) {
  if (n < k) {
    return 0;
  }
  if (k <= 0) {
    return 0;
  }
  return getC(n - 1, k - 1);
}
Mint quickMod(Mint a, LL b) {
  Mint res = 1;
  while (b > 0) {
    if (b % 2) {
      res *= a;
    }
    a *= a;
    b /= 2;
  }
  return res;
}
template <class TypeName> vector<TypeName> getVector1D(int n, TypeName val) {
  return vector<TypeName>(n, val);
}
template <class TypeName>
vector<vector<TypeName>> getVector2D(int n1, int n2, TypeName val) {
  vector<vector<TypeName>> res(n1);
  for (int i = 0; i < n1; ++i) {
    res[i] = vector<TypeName>(n2, val);
  }
  return res;
}
template <class TypeName>
vector<vector<vector<TypeName>>> getVector3D(int n1, int n2, int n3,
                                             TypeName val) {
  vector<vector<vector<TypeName>>> res(n1);
  for (int i = 0; i < n1; ++i) {
    res[i] = vector<vector<TypeName>>(n2);
    for (int j = 0; j < n2; ++j) {
      res[i][j] = vector<TypeName>(n3, val);
    }
  }
  return res;
}
template <class TypeName>
vector<vector<vector<vector<TypeName>>>> getVector4D(int n1, int n2, int n3,
                                                     int n4, TypeName val) {
  vector<vector<vector<vector<TypeName>>>> res(n1);
  for (int i = 0; i < n1; ++i) {
    res[i] = vector<vector<vector<TypeName>>>(n2);
    for (int j = 0; j < n2; ++j) {
      res[i][j] = vector<vector<TypeName>>(n3);
      for (int k = 0; k < n3; ++k) {
        res[i][j][k] = vector<TypeName>(n4, val);
      }
    }
  }
  return res;
}

LL Mint::md = 1000000007;
LL md = 1000000007;

class Mat {
public:
  LL n;
  vector<vector<LL>> a;
  Mat(LL _n, bool makeIdentity = false) : n(_n) {
    a = getVector2D(n, n, (LL)(0));
    if (makeIdentity) {
      for (int i = 0; i < n; ++i) {
        a[i][i] = 1;
      }
    }
  }
  Mat(const Mat &mat) {
    n = mat.n;
    a = mat.a;
  }
  void square() {
    auto tmp = getVector2D(n, n, (LL)(0));
    for (int i = 0; i < n; ++i) {
      for (int j = 0; j < n; ++j) {
        for (int k = 0; k < n; ++k) {
          tmp[i][j] += a[i][k] * a[k][j];
          tmp[i][j] = mod(tmp[i][j], md);
        }
      }
    }
    // for (int i = 0; i < n; ++i) {
    //     for (int j = 0; j < n; ++j) {
    //         tmp[i][j] = mod(tmp[i][j], md);
    //     }
    // }
    swap(a, tmp);
    tmp.clear();
  }
  void multi(Mat &b) {
    auto tmp = getVector2D(n, n, (LL)(0));
    for (int i = 0; i < n; ++i) {
      for (int j = 0; j < n; ++j) {
        for (int k = 0; k < n; ++k) {
          tmp[i][j] += a[i][k] * b.a[k][j];
          tmp[i][j] = mod(tmp[i][j], md);
        }
      }
    }
    // for (int i = 0; i < n; ++i) {
    //     for (int j = 0; j < n; ++j) {
    //         tmp[i][j] = mod(tmp[i][j], md);
    //     }
    // }
    swap(a, tmp);
    tmp.clear();
  }
  vector<LL> vectorMulti(vector<LL> &x) {
    vector<LL> res(n, 0);
    for (int i = 0; i < n; ++i) {
      for (int j = 0; j < n; ++j) {
        res[i] += x[j] * a[j][i];
        res[i] = mod(res[i], md);
      }
      // res[i] = mod(res[i], md);
    }
    return res;
  }

  ~Mat() { a.clear(); }
};

Mat MatToOrderN(Mat &a, LL n) {
  Mat aa = a;
  Mat res(a.n, true);
  while (n > 0) {
    if (n & 1) {
      res.multi(aa);
    }
    aa.square();
    n /= 2;
  }
  return res;
}

LL work() {
  LL n;
  cin >> n;
  if (n < 5) {
    return 0;
  }

  Mat a(32, false);
  for (int jNew = 0; jNew <= 15; ++jNew) {
    for (int pNew = 0; pNew < 2; ++pNew) {
      int y = jNew * 2 + pNew;
      int jOld = jNew, pOld = pNew ^ 1;
      int x = jOld * 2 + pOld;
      a.a[x][y] = 1;

      if (pNew == 0 && jNew > 0) {
        for (jOld = 0; jOld < jNew; ++jOld) {
          x = jOld * 2 + 1;
          a.a[x][y] = 1;
          if (jOld > 4) {
            x = jOld * 2;
            a.a[x][y] = 1;
          }
        }
      }
    }
  }
  // for (int i = 0; i < 32; ++i) {
  //     for (int j = 0; j < 32; ++j) {
  //         cout << a.a[i][j] << ' ';
  //     }
  //     cout << endl;
  // }
  a = MatToOrderN(a, n - 5);
  vector<LL> res(32, 0);
  for (int i = 0; i <= 15; ++i) {
    res[i * 2] = 1;
  }
  res = a.vectorMulti(res);
  LL ans = mod(res[30] + res[31], md);
  return ans;
}

int main() {
  int T;
  cin >> T;
  for (int i = 0; i < T; ++i) {
    cout << work() << endl;
  }
  return 0;
}