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

#ifdef LOCAL
#define _GLIBCXX_DEBUG
#endif

#include <bits/stdc++.h>

using namespace std;

using ll = long long;
using ull = unsigned long long;
using ld = long double;

mt19937 mt(736);

const int mod = int(1e9) + 7;

void add(int &a, int b) {
  a += b;
  if (a >= mod)
    a -= mod;
}

void sub(int &a, int b) {
  a -= b;
  if (a < 0)
    a += mod;
}

int mult(const int a, const int b) { return (ll)a * b % mod; }

int mult(const int a, const int b, const int c) { return mult(mult(a, b), c); }

struct reset_arr {
  vector<int> arr, last;
  int timer;

  reset_arr(const int n) : arr(n, 0), last(n, 0), timer(1) {}

  void reset() { timer++; }

  //	const int &operator[](const int i) const
  //	{
  //		if (last[i] < timer)
  //		{
  //			last[i] = timer;
  //			arr[i] = 0;
  //		}
  //		return arr[i];
  //	}

  int &operator[](const int i) {
    if (last[i] < timer) {
      last[i] = timer;
      arr[i] = 0;
    }
    return arr[i];
  }
};

void solve(istream &cin = std::cin, ostream &cout = std::cout) {
  int bits;
  cin >> bits;
  --bits;

  vector<int> p(1 << bits);
  for (auto &it : p) {
    cin >> it;
    --it;
  }

  vector<int> inv(1 << (bits + 1), 1);
  for (int i = 2; i < int(inv.size()); i++) {
    inv[i] = mod - mult(mod / i, inv[mod % i]);
    assert(mult(i, inv[i]) == 1);
  }

  vector<int> prod_to_root((1 << bits) * 2, 1);
  auto inv_prod = prod_to_root;
  for (int i = 2; i < int(prod_to_root.size()); i++) {
    prod_to_root[i] = mult(prod_to_root[i / 2], i);
    inv_prod[i] = mult(inv_prod[i / 2], inv[i]);
  }

  reset_arr black_sum(prod_to_root.size()), white_sum(prod_to_root.size());

  int ans = 0;
  for (int fst_diff = 0; fst_diff < bits; fst_diff++) {
    for (int mask = 0; mask < (1 << fst_diff); mask++) {
      const int lca = (1 << fst_diff) + mask;
      const int scale = (1 << (bits - fst_diff - 1));

      black_sum.reset();
      white_sum.reset();

      static vector<int> vs;
      vs.clear();

      for (int i = 2 * mask * scale; i < (2 * mask + 2) * scale; i++) {
        const int our = i + (1 << bits);
        const int target = p[i] + (1 << bits);
        vs.emplace_back(target);

        if (i >= (2 * mask + 1) * scale)
          add(black_sum[target],
              mult(prod_to_root[our], inv_prod[lca], target));
        else
          add(white_sum[target],
              mult(prod_to_root[our], inv_prod[lca / 2], target));
      }

      sort(vs.begin(), vs.end());

      auto recalc = [&](const int v) {
        for (int son = 2 * v; son <= 2 * v + 1; son++) {
          add(black_sum[v], mult(black_sum[son], v));
          add(white_sum[v], mult(white_sum[son], v));
          add(ans, mult(black_sum[son], v, white_sum[son ^ 1]));
        }
      };

      while (vs.size() > 1U) {
        for (auto &it : vs)
          it /= 2;
        vs.resize(unique(vs.begin(), vs.end()) - vs.begin());

        for (auto it : vs)
          recalc(it);
      }
    }
  }

  cout << ans << endl;

  //	exit(0);
}

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);

  cout << fixed;

#ifdef LOCAL
  auto st = clock();

  ifstream fin("../input.txt");

  do {
    solve(fin);

    cout << "===" << endl;

    string str;
    while (getline(fin, str) && str != string(max(1, (int)str.size()), '='))
      ;
  } while (fin);

  cout << setprecision(6) << "clock: " << ld(clock() - st) / CLOCKS_PER_SEC
       << endl;
#else
  solve();
#endif

  return 0;
}


Error: TLE

Fixed submission:

#ifdef LOCAL
#define _GLIBCXX_DEBUG
#endif

#include <bits/stdc++.h>

using namespace std;

using ll = long long;
using ull = unsigned long long;
using ld = long double;

mt19937 mt(736);

const int mod = int(1e9) + 7;

void add(int &a, int b) {
  a += b;
  if (a >= mod)
    a -= mod;
}

void sub(int &a, int b) {
  a -= b;
  if (a < 0)
    a += mod;
}

int mult(const int a, const int b) { return (ll)a * b % mod; }

int mult(const int a, const int b, const int c) { return mult(mult(a, b), c); }

struct reset_arr {
  vector<int> arr, last;
  int timer;

  reset_arr(const int n) : arr(n, 0), last(n, 0), timer(1) {}

  void reset() { timer++; }

  //	const int &operator[](const int i) const
  //	{
  //		if (last[i] < timer)
  //		{
  //			last[i] = timer;
  //			arr[i] = 0;
  //		}
  //		return arr[i];
  //	}

  int &operator[](const int i) {
    if (last[i] < timer) {
      last[i] = timer;
      arr[i] = 0;
    }
    return arr[i];
  }
};

void solve(istream &cin = std::cin, ostream &cout = std::cout) {
  int bits;
  cin >> bits;
  --bits;

  vector<int> p(1 << bits);
  for (auto &it : p) {
    cin >> it;
    --it;
  }

  vector<int> inv(1 << (bits + 1), 1);
  for (int i = 2; i < int(inv.size()); i++) {
    inv[i] = mod - mult(mod / i, inv[mod % i]);
    assert(mult(i, inv[i]) == 1);
  }

  vector<int> prod_to_root((1 << bits) * 2, 1);
  auto inv_prod = prod_to_root;
  for (int i = 2; i < int(prod_to_root.size()); i++) {
    prod_to_root[i] = mult(prod_to_root[i / 2], i);
    inv_prod[i] = mult(inv_prod[i / 2], inv[i]);
  }

  reset_arr black_sum(prod_to_root.size()), white_sum(prod_to_root.size());

  int ans = 0;
  for (int fst_diff = 0; fst_diff < bits; fst_diff++) {
    for (int mask = 0; mask < (1 << fst_diff); mask++) {
      const int lca = (1 << fst_diff) + mask;
      const int scale = (1 << (bits - fst_diff - 1));

      black_sum.reset();
      white_sum.reset();

      vector<int> vs;
      for (int i = 2 * mask * scale; i < (2 * mask + 2) * scale; i++) {
        const int our = i + (1 << bits);
        const int target = p[i] + (1 << bits);
        vs.emplace_back(target);

        if (i >= (2 * mask + 1) * scale)
          add(black_sum[target],
              mult(prod_to_root[our], inv_prod[lca], target));
        else
          add(white_sum[target],
              mult(prod_to_root[our], inv_prod[lca / 2], target));
      }
      sort(vs.begin(), vs.end());

      auto recalc = [&](const int v) {
        for (int son = 2 * v; son <= 2 * v + 1; son++) {
          add(black_sum[v], mult(black_sum[son], v));
          add(white_sum[v], mult(white_sum[son], v));
          add(ans, mult(black_sum[son], v, white_sum[son ^ 1]));
        }
      };

      for (int i = int(inv_prod.size()) / 2 - 1; i >= 1; i--)
        recalc(i);

      //			while (vs.size() > 1U)
      //			{
      //				for (auto &it : vs)
      //					it /= 2;
      //				vs.resize(unique(vs.begin(), vs.end()) -
      //vs.begin());
      //
      //				for (auto it : vs)
      //					recalc(it);
      //			}
    }
  }

  cout << ans << endl;

  //	exit(0);
}

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);

  cout << fixed;

#ifdef LOCAL
  auto st = clock();

  ifstream fin("../input.txt");

  do {
    solve(fin);

    cout << "===" << endl;

    string str;
    while (getline(fin, str) && str != string(max(1, (int)str.size()), '='))
      ;
  } while (fin);

  cout << setprecision(6) << "clock: " << ld(clock() - st) / CLOCKS_PER_SEC
       << endl;
#else
  solve();
#endif

  return 0;
}