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

#include <algorithm>
#include <bitset>
#include <cassert>
#include <chrono>
#include <cmath>
#include <complex>
#include <cstring>
#include <ctime>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
using namespace std;
#define int long long
#define endl '\\n'
typedef long long ll;

const int N = (1 << 18) + 10, mod = 1e9 + 7;
int h;
int perm[N], saveVal[N];
vector<int> nodesInSaveVal;

void dfs1(int node, int pathProd) {
  pathProd = (pathProd * node) % mod;
  if (node >= (1 << (h - 1))) { // node is leaf
    int nodeOther = perm[node];
    while (true) {
      pathProd *= nodeOther;
      pathProd %= mod;
      nodesInSaveVal.push_back(nodeOther);
      saveVal[nodeOther] += pathProd;
      saveVal[nodeOther] %= mod;
      if (nodeOther == 1)
        break;
      nodeOther /= 2;
    }
  } else {
    dfs1(2 * node, pathProd);
    dfs1(2 * node + 1, pathProd);
  }
}

int currSum;

void dfs2(int node, int pathProd) {
  pathProd = (pathProd * node) % mod;
  if (node >= (1 << (h - 1))) { // node is leaf
    int nodeOther = perm[node];
    int prevNode = nodeOther;
    while (true) {
      pathProd *= nodeOther;
      pathProd %= mod;

      if (prevNode > nodeOther) {
        currSum += (pathProd * saveVal[prevNode ^ 1]) % mod;
      }
      currSum %= mod;
      if (nodeOther == 1)
        break;
      prevNode = nodeOther;
      nodeOther /= 2;
    }
  } else {
    dfs2(2 * node, pathProd);
    dfs2(2 * node + 1, pathProd);
  }
}

signed main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);
  cin >> h;
  for (int i = (1 << (h - 1)); i < (1 << h); ++i) {
    cin >> perm[i];
    perm[i] += (1 << (h - 1)) - 1;
  }
  int res = 0;
  for (int lca = 1; lca < (1 << (h - 1)); ++lca) {
    int L = 2 * lca;
    int R = 2 * lca + 1;
    nodesInSaveVal.clear();
    dfs1(L, 1);
    currSum = 0;
    dfs2(R, 1);
    res += (lca * currSum) % mod;
    res %= mod;
    for (int node : nodesInSaveVal) {
      saveVal[node] = 0;
    }
  }
  cout << res;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <bitset>
#include <cassert>
#include <chrono>
#include <cmath>
#include <complex>
#include <cstring>
#include <ctime>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
using namespace std;
#define int long long
#define endl '\\n'
typedef long long ll;

const int N = (1 << 18) + 10, mod = 1e9 + 7;
int h;
int perm[N], height[N], saveVal[N];
vector<int> nodesInSaveVal;

void dfs1(int node, int pathProd) {
  pathProd = (pathProd * node) % mod;
  if (node >= (1 << (h - 1))) { // node is leaf
    int nodeOther = perm[node];
    while (true) {
      pathProd *= nodeOther;
      pathProd %= mod;
      nodesInSaveVal.push_back(nodeOther);
      saveVal[nodeOther] += pathProd;
      saveVal[nodeOther] %= mod;
      if (nodeOther == 1)
        break;
      nodeOther /= 2;
    }
  } else {
    dfs1(2 * node, pathProd);
    dfs1(2 * node + 1, pathProd);
  }
}

int currSum;

void dfs2(int node, int pathProd) {
  pathProd = (pathProd * node) % mod;
  if (node >= (1 << (h - 1))) { // node is leaf
    int nodeOther = perm[node];
    int prevNode = nodeOther;
    while (true) {
      pathProd *= nodeOther;
      pathProd %= mod;

      if (prevNode > nodeOther) {
        currSum += (pathProd * saveVal[prevNode ^ 1]) % mod;
      }
      currSum %= mod;
      if (nodeOther == 1)
        break;
      prevNode = nodeOther;
      nodeOther /= 2;
    }
  } else {
    dfs2(2 * node, pathProd);
    dfs2(2 * node + 1, pathProd);
  }
}

signed main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);
  cin >> h;
  for (int i = (1 << (h - 1)); i < (1 << h); ++i) {
    cin >> perm[i];
    perm[i] += (1 << (h - 1)) - 1;
  }
  height[1] = 1;
  for (int i = 1; i < (1 << (h - 1)); ++i) {
    height[2 * i] = height[2 * i + 1] = 1 + height[i];
  }
  int res = 0;
  for (int lca = 1; lca < (1 << (h - 1)); ++lca) {
    int L = 2 * lca;
    int R = 2 * lca + 1;
    nodesInSaveVal.clear();
    dfs1(L, 1);
    currSum = 0;
    dfs2(R, 1);
    res += (lca * currSum) % mod;
    res %= mod;
    for (int node : nodesInSaveVal) {
      saveVal[node] = 0;
    }
    for (int i = 0; i < N; ++i)
      assert(saveVal[i] == 0);
  }
  cout << res;
}