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

#include <bits/stdc++.h>
using namespace std;
#define F first
#define S second
#define PB push_back
#define MP make_pair
#define lc 2 * v
#define rc 2 * v + 1
#define mid (s + e) / 2
#define ll long long
#define int long long
#define ld long double
#define pii pair<int, int>
#define pll pair<long long, long long>
#define FAST                                                                   \\
  ios::sync_with_stdio(false);                                                 \\
  cin.tie(0);
#pragma GCC optimize("Ofast")
#pragma GCC target("sse,sse2,sse3,ssse3,sse4,avx,avx2")

const ll maxn = 3e5 + 100, N = 2e6 + 10, SQ = 320, base = 97, mod = 1e9 + 7,
         INF = 1e18, lg = 19;

int H, S, E, ans, p[maxn], pos[maxn], sum[maxn];
vector<int> leaf;
bool mark[maxn];

inline void get(int v) {
  int s = v, e = v;
  for (int i = 0; i < H - 1; ++i) {
    s = (s * 2);
    e = (e * 2) + 1;
  }
  S = s;
  E = e;
  for (int i = s; i <= e; ++i) {
    leaf.PB(i);
    mark[i] = 1;
    pos[i] = i - s;
  }
}

inline void undo_root(int v) {
  while (v >= 1) {
    sum[v] = 0;
    v /= 2;
  }
}
inline void undo(int v) {
  int s = v, e = v;
  while (!mark[s] && !mark[e]) {
    s = (2 * s);
    e = (2 * e + 1);
  }
  for (int i = s; i <= e; ++i) {
    int idx = pos[i];
    undo_root(leaf[p[idx]]);
  }
}

inline void go_root1(int v, int val) {
  while (v >= 1) {
    val = (val * v) % mod;
    sum[v] += val;
    v /= 2;
  }
}
inline void go_root2(int v, int val) {
  while (v > 1) {
    val = (val * v) % mod;
    int par = v / 2;
    int other_child = (par * 2 == v ? v + 1 : v - 1);
    ans = (ans + (((par * sum[other_child]) % mod) * val) % mod) % mod;
    v /= 2;
  }
}

void go_l(int v, int val) {
  val = (val * v) % mod;
  if (mark[v]) {
    int idx = pos[v];
    go_root1(leaf[p[idx]], val);
    return;
  }
  go_l(2 * v, val);
  go_l(2 * v + 1, val);
}
void go_r(int v, int val) {
  val = (val * v) % mod;
  if (mark[v]) {
    int idx = pos[v];
    // cout << leaf[idx] << ' ' << leaf[p[idx]] << endl;
    go_root2(leaf[p[idx]], val);
    return;
  }
  go_r(2 * v, val);
  go_r(2 * v + 1, val);
}

void fix_lca(int v) {
  go_l(2 * v, v);
  go_r(2 * v + 1, 1);
  undo(2 * v);
}

int32_t main() {
  FAST cin >> H;
  for (int i = 0; i < (1 << (H - 1)); ++i) {
    cin >> p[i];
    p[i]--;
  }
  get(1);

  for (int i = 1; i < leaf[0]; ++i) {
    fix_lca(i);
  }
  cout << ans << '\\n';
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define F first
#define S second
#define PB push_back
#define MP make_pair
#define lc 2 * v
#define rc 2 * v + 1
#define mid (s + e) / 2
#define ll long long
#define int long long
#define ld long double
#define pii pair<int, int>
#define pll pair<long long, long long>
#define FAST                                                                   \\
  ios::sync_with_stdio(false);                                                 \\
  cin.tie(0);
#pragma GCC optimize("Ofast")
#pragma GCC target("sse,sse2,sse3,ssse3,sse4,avx,avx2")

const ll maxn = 3e5 + 100, N = 2e6 + 10, SQ = 320, base = 97, mod = 1e9 + 7,
         INF = 1e18, lg = 19;

int H, S, E, ans, p[maxn], pos[maxn], sum[maxn];
vector<int> leaf;
bool mark[maxn];

inline void get(int v) {
  int s = v, e = v;
  for (int i = 0; i < H - 1; ++i) {
    s = (s * 2);
    e = (e * 2) + 1;
  }
  S = s;
  E = e;
  for (int i = s; i <= e; ++i) {
    leaf.PB(i);
    mark[i] = 1;
    pos[i] = i - s;
  }
}

inline void go_root1(int v, int val) {
  while (v >= 1) {
    val = (val * v) % mod;
    sum[v] += val;
    v /= 2;
  }
}

inline void go_root2(int v, int val) {
  while (v > 1) {
    val = (val * v) % mod;
    int par = v / 2;
    int other_child = (par * 2 == v ? v + 1 : v - 1);
    ans = (ans + (((par * sum[other_child]) % mod) * val) % mod) % mod;
    v /= 2;
  }
}

void go_l(int v, int val) {
  val = (val * v) % mod;
  if (mark[v]) {
    int idx = pos[v];
    // cout << leaf[idx] << ' ' << leaf[p[idx]] << endl;
    go_root1(leaf[p[idx]], val);
    return;
  }
  go_l(2 * v, val);
  go_l(2 * v + 1, val);
}

void go_r(int v, int val) {
  val = (val * v) % mod;
  if (mark[v]) {
    int idx = pos[v];
    // cout << leaf[idx] << ' ' << leaf[p[idx]] << endl;
    go_root2(leaf[p[idx]], val);
    return;
  }
  go_r(2 * v, val);
  go_r(2 * v + 1, val);
}

void fix_lca(int v) {
  go_l(2 * v, v);
  go_r(2 * v + 1, 1);
  for (int i = 0; i < (1ll << H) + 1; ++i) {
    sum[i] = 0;
  }
}

int32_t main() {
  FAST cin >> H;
  for (int i = 0; i < (1 << (H - 1)); ++i) {
    cin >> p[i];
    p[i]--;
  }
  get(1);
  // for (auto u : leaf) {
  //     cout << u << ' ';
  // }
  // cout << endl;
  for (int i = 1; i < leaf[0]; ++i) {
    fix_lca(i);
  }
  cout << ans << '\\n';
}