Problem p02866 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>

using namespace std;
using ll = long long;
using db = double;
using vll = vector<long long>;

template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
#define che(a, string) cout << "//" << string << "==" << (a) << "//" << endl;
#define SORT(a) sort((a).begin(), (a).end())
#define rSORT(a) reverse((a).begin(), (a).end())
#define For(i, a, b) for (int i = (a); i < (b); ++i)
#define rep(i, n) For(i, 0, n)
#define leftunique(a)                                                          \
  {                                                                            \
    SORT(a);                                                                   \
    (a).erase(unique((a).begin(), (a).end()), (a).end());                      \
  }
// leftuniqueは配列から同じ要素のものを取り除く 例 4 6 2 1 3 4 1 2 なら 1 2 3 4
// 6 になる
bool IsInt(double a) {
  int b = a / 1;
  if (a == b) {
    return true;
  } else {
    return false;
  }
}
bool coY() { cout << "Yes" << endl; }
bool coN() { cout << "No" << endl; }
const int mod = 998244353;
const ll INF = 1LL << 60;

// a^n mod を計算する
long long modpow(long long a, long long n, long long mod) {
  long long res = 1;
  while (n > 0) {
    if (n & 1)
      res = res * a % mod;
    a = a * a % mod;
    n >>= 1;
  }
  return res;
}

// ここから書き始める

int main() {
  cout << setprecision(10);
  ll n;
  cin >> n;
  map<ll, ll> tree;
  rep(i, n) {
    ll a;
    cin >> a; // 入力で距離を与える
    if (i == 0) {
      // 距離は0じゃないといけない!!
      if (a != 0) {
        cout << 0 << endl; // 1個目の距離が0じゃない木は存在しない
        return 0;
      } else {
        tree[a] = 1;
      }
    }
    // i != 0 の ときは map にどんどんぶち込んでいく。
    else {
      if (tree.count(a)) {
        tree[a]++;
      } else
        tree[a] = 1;
    }
  }
  // 一回マップの中を見てみる
  ll ans = 1;
  rep(i, n) {
    if (i == 0) {
      if (tree[i] != 1) { // 0は1個しかダメ。
        cout << 0 << endl;
        return 0;
      }
    } else if (tree[i] != 0) {
      ans *= modpow(tree[i - 1], tree[i], mod);
      ans %= mod;
    } else {
      For(j, i + 1, n) {
        if (tree[j] != 0) {
          cout << 0 << endl;
          return 0;
        }
      }
    }
  }
  cout << ans << endl;
}


Fixed submission:
#include <bits/stdc++.h>

using namespace std;
using ll = long long;
using db = double;
using vll = vector<long long>;

template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
#define che(a, string) cout << "//" << string << "==" << (a) << "//" << endl;
#define SORT(a) sort((a).begin(), (a).end())
#define rSORT(a) reverse((a).begin(), (a).end())
#define For(i, a, b) for (int i = (a); i < (b); ++i)
#define rep(i, n) For(i, 0, n)
#define leftunique(a)                                                          \
  {                                                                            \
    SORT(a);                                                                   \
    (a).erase(unique((a).begin(), (a).end()), (a).end());                      \
  }
// leftuniqueは配列から同じ要素のものを取り除く 例 4 6 2 1 3 4 1 2 なら 1 2 3 4
// 6 になる
bool IsInt(double a) {
  int b = a / 1;
  if (a == b) {
    return true;
  } else {
    return false;
  }
}
bool coY() { cout << "Yes" << endl; }
bool coN() { cout << "No" << endl; }
const int mod = 998244353;
const ll INF = 1LL << 60;

// a^n mod を計算する
long long modpow(long long a, long long n, long long mod) {
  long long res = 1;
  while (n > 0) {
    if (n & 1)
      res = res * a % mod;
    a = a * a % mod;
    n >>= 1;
  }
  return res;
}

// ここから書き始める

int main() {
  cout << setprecision(10);
  ll n;
  cin >> n;
  map<ll, ll> tree;
  rep(i, n) {
    ll a;
    cin >> a; // 入力で距離を与える
    if (i == 0) {
      // 距離は0じゃないといけない!!
      if (a != 0) {
        cout << 0 << endl; // 1個目の距離が0じゃない木は存在しない
        return 0;
      } else {
        tree[a] = 1;
      }
    }
    // i != 0 の ときは map にどんどんぶち込んでいく。
    else {
      if (tree.count(a)) {
        tree[a]++;
      } else
        tree[a] = 1;
    }
  }
  // 一回マップの中を見てみる
  ll ans = 1;
  rep(i, n) {
    if (i == 0) {
      if (tree[i] != 1) { // 0は1個しかダメ。
        cout << 0 << endl;
        return 0;
      }
    } else if (tree[i] != 0) {
      ans *= modpow(tree[i - 1], tree[i], mod);
      ans %= mod;
    } else {
      For(j, i + 1, n) {
        if (tree[j] != 0) {
          cout << 0 << endl;
          return 0;
        }
      }
      break;
    }
  }
  cout << ans << endl;
}