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

#include <bits/stdc++.h>
using namespace std;
const int N = 3e5 + 100;
const int P = 1e9 + 7;
#define fr(i, a, b) for (int i = a; i <= b; ++i)
#define nfr(i, a, b) for (int i = a; i >= b; --i)

inline int pow_mod(int x, int k) {
  int res = 1, base = x;
  while (k) {
    if (k & 1)
      res = 1LL * res * base % P;
    k >>= 1;
    base = 1LL * base * base % P;
  }
  return res;
}

inline void add(int &x, int y) {
  x += y;
  if (x >= P)
    x -= P;
}

struct segment_tree {
  int val[N * 20], ch[N * 20][2];
  int cnt;
#define lc ch[x][0]
#define rc ch[x][1]
#define mid (l + r) / 2
  inline int newnode() {
    ++cnt;
    ch[cnt][0] = ch[cnt][1] = val[cnt] = 0;
    return cnt;
  }
  inline void modify(int &x, int l, int r, int pos, int v) {
    if (!x)
      x = newnode();
    add(val[x], v);
    if (l == r)
      return;
    if (pos <= mid)
      modify(lc, l, mid, pos, v);
    else
      modify(rc, mid + 1, r, pos, v);
  }
  inline int merge(int x, int y, int l, int r) {
    if (x * y == 0)
      return x | y;
    int pre = x;
    x = newnode();
    add(val[x] = val[pre], val[y]);
    if (l == r)
      return x;
    lc = merge(ch[pre][0], ch[y][0], l, mid);
    rc = merge(ch[pre][1], ch[y][1], mid + 1, r);
    return x;
  }
  inline int query(int x, int l, int r, int ql, int qr) {
    if (ql > r || qr < l)
      return 0;
    if (ql <= l && r <= qr) {
      return val[x];
    }
    int res = 0;
    add(res = query(lc, l, mid, ql, qr), query(rc, mid + 1, r, ql, qr));
    return res;
  }
  inline int qid(int x, int l, int r, int ql, int qr) {
    if (ql > r || qr < l)
      return 0;
    if (ql <= l && r <= qr) {
      return x;
    }
    return qid(lc, l, mid, ql, qr) + qid(rc, mid + 1, r, ql, qr);
  }
} s;

int h, l[N], n;
int rt[N], L[N], R[N];
int prod[N], inv[N];
int lb, rb;

void init() {
  scanf("%d", &h);
  n = (1 << h);

  fr(i, 1, n - 1) rt[i] = s.newnode();
  lb = n / 2, rb = n - 1;
  fr(i, n / 2, n - 1) {
    scanf("%d", l + i);
    l[i] += n / 2 - 1;
  }

  prod[1] = inv[1] = 1;
  fr(i, 2, n - 1) {
    prod[i] = 1LL * prod[i / 2] * i % P;
    inv[i] = pow_mod(prod[i], P - 2);
    inv[i] = 1LL * inv[i] * inv[i] % P;
  }

  nfr(i, n - 1, 1) {
    if (i >= n / 2) {
      s.modify(rt[i], lb, rb, l[i], 1LL * prod[i] * prod[l[i]] % P);
      L[i] = R[i] = i;
    } else {
      rt[i] = s.merge(rt[i << 1], rt[i << 1 | 1], lb, rb);
      L[i] = L[i << 1];
      R[i] = R[i << 1 | 1];
    }
  }

  int ans = 0;
  fr(i, n / 2, n - 1) {
    for (int j = i; j > 1; j /= 2) {
      int cj = (j == (j / 2) * 2) ? j + 1 : j - 1;

      int k = 1;
      int kid = rt[cj];
      while (k != l[i]) {
        int ck = 0;
        if (l[i] >= L[k << 1 | 1])
          ck = 1;
        int qv = s.val[s.ch[kid][ck ^ 1]];
        add(ans, 1LL * qv * inv[j / 2] % P * inv[k] % P * prod[i] % P *
                     prod[l[i]] % P * (j / 2) % P * k % P);
        k = (k << 1 | ck);
        kid = s.ch[kid][ck];
      }
    }
  }
  printf("%d\\n", 1LL * ans * pow_mod(2, P - 2) % P);
}

int main() {
  init();
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
const int N = 3e5 + 100;
const int P = 1e9 + 7;
#define fr(i, a, b) for (int i = a; i <= b; ++i)
#define nfr(i, a, b) for (int i = a; i >= b; --i)

inline int pow_mod(int x, int k) {
  int res = 1, base = x;
  while (k) {
    if (k & 1)
      res = 1LL * res * base % P;
    k >>= 1;
    base = 1LL * base * base % P;
  }
  return res;
}

inline void add(int &x, int y) {
  x += y;
  if (x >= P)
    x -= P;
}

struct segment_tree {
  int val[N * 20], ch[N * 20][2];
  int cnt;
#define lc ch[x][0]
#define rc ch[x][1]
#define mid (l + r) / 2
  inline int newnode() {
    ++cnt;
    ch[cnt][0] = ch[cnt][1] = val[cnt] = 0;
    return cnt;
  }
  inline void modify(int &x, int l, int r, int pos, int v) {
    if (!x)
      x = newnode();
    add(val[x], v);
    if (l == r)
      return;
    if (pos <= mid)
      modify(lc, l, mid, pos, v);
    else
      modify(rc, mid + 1, r, pos, v);
  }
  inline int merge(int x, int y, int l, int r) {
    if (x * y == 0)
      return x | y;
    int pre = x;
    x = newnode();
    add(val[x] = val[pre], val[y]);
    if (l == r)
      return x;
    lc = merge(ch[pre][0], ch[y][0], l, mid);
    rc = merge(ch[pre][1], ch[y][1], mid + 1, r);
    return x;
  }
  inline int query(int x, int l, int r, int ql, int qr) {
    if (ql > r || qr < l)
      return 0;
    if (ql <= l && r <= qr) {
      return val[x];
    }
    int res = 0;
    add(res = query(lc, l, mid, ql, qr), query(rc, mid + 1, r, ql, qr));
    return res;
  }
} s;

int h, l[N], n;
int rt[N], L[N], R[N];
int prod[N], inv[N];
int lb, rb;

void init() {
  scanf("%d", &h);
  n = (1 << h);

  fr(i, 1, n - 1) rt[i] = s.newnode();
  lb = n / 2, rb = n - 1;
  fr(i, n / 2, n - 1) {
    scanf("%d", l + i);
    l[i] += n / 2 - 1;
  }

  prod[1] = inv[1] = 1;
  fr(i, 2, n - 1) {
    prod[i] = 1LL * prod[i / 2] * i % P;
    inv[i] = pow_mod(prod[i], P - 2);
    inv[i] = 1LL * inv[i] * inv[i] % P;
  }

  nfr(i, n - 1, 1) {
    if (i >= n / 2) {
      s.modify(rt[i], lb, rb, l[i], 1LL * prod[i] * prod[l[i]] % P);
      cerr << l[i] << ' ' << 1LL * prod[i] * prod[l[i]] % P << endl;
      L[i] = R[i] = i;
    } else {
      rt[i] = s.merge(rt[i << 1], rt[i << 1 | 1], lb, rb);
      L[i] = L[i << 1];
      R[i] = R[i << 1 | 1];
    }
  }

  int ans = 0;
  fr(i, n / 2, n - 1) {
    for (int j = i; j > 1; j /= 2) {
      int cj = (j == (j / 2) * 2) ? j + 1 : j - 1;
      for (int k = l[i]; k > 1; k /= 2) {
        int ck = (k == (k / 2) * 2) ? k + 1 : k - 1;
        int qv = s.query(rt[cj], lb, rb, L[ck], R[ck]);
        add(ans, 1LL * qv * inv[j / 2] % P * inv[k / 2] % P * prod[i] % P *
                     prod[l[i]] % P * (j / 2) % P * (k / 2) % P);
      }
    }
  }
  printf("%d\\n", 1LL * ans * pow_mod(2, P - 2) % P);
}

int main() {
  init();
  return 0;
}