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

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#pragma GCC target("avx2")
#pragma GCC optimization("O3")
#pragma GCC optimization("unroll-loops")

using namespace std;
using namespace __gnu_pbds;

#define fi first
#define se second
#define mp make_pair
#define pb push_back
#define fbo find_by_order
#define ook order_of_key

typedef long long ll;
typedef pair<int, int> ii;
typedef vector<int> vi;
typedef long double ld;
typedef tree<ll, null_type, less<ll>, rb_tree_tag,
             tree_order_statistics_node_update>
    pbds;

vector<int> fact;
vector<int> ifact;
vector<int> inv;
vector<int> pow2;
const int MOD = (1e9 + 7);
int add(int a, int b) {
  a += b;
  while (a >= MOD)
    a -= MOD;
  return a;
}
void radd(int &a, int b) { a = add(a, b); }
int mult(int a, int b) { return (a * 1LL * b) % MOD; }
void rmult(int &a, int b) { a = mult(a, b); }
int modpow(int a, int b) {
  int r = 1;
  while (b) {
    if (b & 1)
      r = mult(r, a);
    a = mult(a, a);
    b >>= 1;
  }
  return r;
}
int choose(int a, int b) {
  if (a < b)
    return 0;
  if (b == 0)
    return 1;
  if (a == b)
    return 1;
  return mult(fact[a], mult(ifact[b], ifact[a - b]));
}
int inverse(int a) { return modpow(a, MOD - 2); }
void init(int _n) {
  fact.clear();
  ifact.clear();
  inv.clear();
  pow2.clear();
  fact.resize(_n + 1);
  ifact.resize(_n + 1);
  inv.resize(_n + 1);
  pow2.resize(_n + 1);
  pow2[0] = 1;
  ifact[0] = 1;
  fact[0] = 1;
  for (int i = 1; i <= _n; i++) {
    pow2[i] = add(pow2[i - 1], pow2[i - 1]);
    fact[i] = mult(fact[i - 1], i);
    // ifact[i]=mult(ifact[i-1],inv[i]);
  }
  ifact[_n] = inverse(fact[_n]);
  for (int i = _n - 1; i >= 1; i--) {
    ifact[i] = mult(ifact[i + 1], i + 1);
  }
  for (int i = 1; i <= _n; i++) {
    inv[i] = mult(fact[i - 1], ifact[i]);
  }
}

ll hsh(ll a, ll b) { return (a << 20) + b; }

int P[(1 << 18) + 20];
int lc[(1 << 18) + 20];
int a[(1 << 18) + 20];
int F[(1 << 18) + 20];
int h;
int N;
int ans;
int n;

vi clrlist;

void mark(int id, int treeno, int prod = 1) {
  if (id == 0)
    return;
  prod = mult(prod, id);
  if (treeno == 1)
    clrlist.pb(id);
  if (id >= N && (!treeno)) {
    mark(a[id], 1, prod);
    return;
  }
  if (treeno == 0) {
    mark(id * 2, treeno, prod);
    mark(id * 2 + 1, treeno, prod);
  } else {
    radd(F[id], prod);
    mark(id >> 1, treeno, prod);
  }
}

void go(int id, int treeno, int prod = 1) {
  if (id == 0)
    return;
  prod = mult(prod, id);
  if (id >= N && (!treeno)) {
    go(a[id], 1, prod);
    return;
  }
  if (treeno == 0) {
    go(id * 2, treeno, prod);
    go(id * 2 + 1, treeno, prod);
  } else {
    if (id > 1) {
      radd(ans, mult(mult(prod, (id >> 1)), F[id ^ 1]));
    }
    go(id >> 1, treeno, prod);
  }
}

void solve(int id) {
  // cerr<<id<<'\\n';
  mark(id * 2, 0, id);
  // cerr<<"M: "<<id<<'\\n';
  go(id * 2 + 1, 0, 1);
  // cerr<<"E: "<<id<<'\\n';
}

int main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  // fix x, p_x, lc_x, lcp_x
  cin >> h;
  N = (1 << (h - 1));
  for (int i = (1 << (h - 1)); i < (1 << h); i++) {
    cin >> a[i];
    a[i] += (1 << (h - 1)) - 1;
  }
  n = (1 << h) - 1;
  for (int i = 1; i <= n; i++) {
    int cur = 1;
    int id = i;
    while (id >= 1) {
      cur = mult(cur, id);
      id >>= 1;
    }
    P[i] = cur;
  }
  // unordered_map<ll,int> ma;
  for (int i = 1; i <= n; i++) {
    int cur = i;
    int id = i / 2;
    while (id >= 1) {
      cur = mult(cur, mult(id, id));
      id >>= 1;
    }
    lc[i] = inverse(cur);
  }
  for (int i = 1; i < N; i++) {
    solve(i);
    for (int x : clrlist)
      F[x] = 0;
    clrlist.clear();
  }
  cout << ans << '\\n';
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace std;
using namespace __gnu_pbds;

#define fi first
#define se second
#define mp make_pair
#define pb push_back
#define fbo find_by_order
#define ook order_of_key

typedef long long ll;
typedef pair<int, int> ii;
typedef vector<int> vi;
typedef long double ld;
typedef tree<ll, null_type, less<ll>, rb_tree_tag,
             tree_order_statistics_node_update>
    pbds;

vector<int> fact;
vector<int> ifact;
vector<int> inv;
vector<int> pow2;
const int MOD = (1e9 + 7);
int add(int a, int b) {
  a += b;
  while (a >= MOD)
    a -= MOD;
  return a;
}
void radd(int &a, int b) { a = add(a, b); }
int mult(int a, int b) { return (a * 1LL * b) % MOD; }
void rmult(int &a, int b) { a = mult(a, b); }
int modpow(int a, int b) {
  int r = 1;
  while (b) {
    if (b & 1)
      r = mult(r, a);
    a = mult(a, a);
    b >>= 1;
  }
  return r;
}
int choose(int a, int b) {
  if (a < b)
    return 0;
  if (b == 0)
    return 1;
  if (a == b)
    return 1;
  return mult(fact[a], mult(ifact[b], ifact[a - b]));
}
int inverse(int a) { return modpow(a, MOD - 2); }
void init(int _n) {
  fact.clear();
  ifact.clear();
  inv.clear();
  pow2.clear();
  fact.resize(_n + 1);
  ifact.resize(_n + 1);
  inv.resize(_n + 1);
  pow2.resize(_n + 1);
  pow2[0] = 1;
  ifact[0] = 1;
  fact[0] = 1;
  for (int i = 1; i <= _n; i++) {
    pow2[i] = add(pow2[i - 1], pow2[i - 1]);
    fact[i] = mult(fact[i - 1], i);
    // ifact[i]=mult(ifact[i-1],inv[i]);
  }
  ifact[_n] = inverse(fact[_n]);
  for (int i = _n - 1; i >= 1; i--) {
    ifact[i] = mult(ifact[i + 1], i + 1);
  }
  for (int i = 1; i <= _n; i++) {
    inv[i] = mult(fact[i - 1], ifact[i]);
  }
}

ll hsh(ll a, ll b) { return (a << 20) + b; }

int P[(1 << 18) + 20];
int lc[(1 << 18) + 20];
int a[(1 << 18) + 20];
// int cnt[(1<<17)*17*17+10];
vector<ii> storage[(1 << 18) + 20];

int main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  init(5);
  // fix x, p_x, lc_x, lcp_x
  int h;
  cin >> h;
  for (int i = (1 << (h - 1)); i < (1 << h); i++) {
    cin >> a[i];
    a[i] += (1 << (h - 1)) - 1;
  }
  int n = (1 << h) - 1;
  for (int i = 1; i <= n; i++) {
    int cur = 1;
    int id = i;
    while (id >= 1) {
      cur = mult(cur, id);
      id >>= 1;
    }
    P[i] = cur;
  }
  // unordered_map<ll,int> ma;
  for (int i = 1; i <= n; i++) {
    int cur = i;
    int id = i / 2;
    while (id >= 1) {
      cur = mult(cur, mult(id, id));
      id >>= 1;
    }
    lc[i] = inverse(cur);
  }
  for (int i = (1 << (h - 1)); i <= n; i++) {
    int j = a[i];
    int cost = mult(P[i], P[j]);
    for (int a = 0; a <= h; a++) {
      for (int b = 0; b <= h; b++) {
        int c = (i >> a);
        int d = (j >> b);
        if (c > 0 && d > 0) {
          storage[c].pb({d, cost});
          // radd(storage[c][d],cost);
          // radd(ma[hsh(c,d)], cost);
        }
      }
    }
  }
  for (int c = 1; c <= n; c++) {
    sort(storage[c].begin(), storage[c].end());
    vector<ii> newst;
    for (int j = int(storage[c].size()) - 1; j >= 0; j--) {
      if (!newst.empty() && newst.back().fi == storage[c][j].fi) {
        radd(newst.back().se, storage[c][j].se);
      } else {
        newst.pb(storage[c][j]);
      }
    }
    reverse(newst.begin(), newst.end());
    storage[c] = newst;
  }
  int ans = 0;
  for (int i = (1 << (h - 1)); i <= n; i++) {
    int j = a[i];
    int cost = mult(P[i], P[j]);
    int res = 0;
    for (int a = 1; a <= h; a++) {
      for (int b = 1; b <= h; b++) {
        int c = (i >> a);
        int d = (j >> b);
        if (c > 0 && d > 0) {
          int c2 = c * 2;
          int d2 = d * 2;
          if ((i >> (a - 1)) == c2)
            c2++;
          if ((j >> (b - 1)) == d2)
            d2++;
          int id =
              lower_bound(storage[c2].begin(), storage[c2].end(), mp(d2, -1)) -
              storage[c2].begin();
          if (id < int(storage[c2].size()) && storage[c2][id].fi == d2) {
            radd(res, mult(storage[c2][id].se, mult(lc[c], lc[d])));
          }
          // radd(res,mult(storage[c2][d2],mult(lc[c],lc[d])));
          // radd(res,mult(ma[hsh(c2,d2)],mult(lc[c],lc[d])));
        }
      }
    }
    res = mult(res, cost);
    radd(ans, res);
  }
  ans = mult(ans, inv[2]);
  cout << ans << '\\n';
}