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

#include <algorithm>
#include <cassert>
#include <chrono>
#include <cmath>
#include <cstdio>
#include <ctime>
#include <functional>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
using namespace std;

namespace {
using Integer = long long; //__int128;
template <class T, class S> istream &operator>>(istream &is, pair<T, S> &p) {
  return is >> p.first >> p.second;
}
template <class T> istream &operator>>(istream &is, vector<T> &vec) {
  for (T &val : vec)
    is >> val;
  return is;
}
template <class T> istream &operator,(istream &is, T &val) { return is >> val; }
template <class T, class S>
ostream &operator<<(ostream &os, const pair<T, S> &p) {
  return os << p.first << " " << p.second;
}
template <class T> ostream &operator<<(ostream &os, const vector<T> &vec) {
  for (size_t i = 0; i < vec.size(); i++)
    os << vec[i] << (i == vec.size() - 1 ? "" : " ");
  return os;
}
template <class T> ostream &operator,(ostream &os, const T &val) {
  return os << " " << val;
}

template <class H> void print(const H &head) { cout << head; }
template <class H, class... T> void print(const H &head, const T &...tail) {
  cout << head << " ";
  print(tail...);
}
template <class... T> void println(const T &...values) {
  print(values...);
  cout << endl;
}

template <class H> void eprint(const H &head) { cerr << head; }
template <class H, class... T> void eprint(const H &head, const T &...tail) {
  cerr << head << " ";
  eprint(tail...);
}
template <class... T> void eprintln(const T &...values) {
  eprint(values...);
  cerr << endl;
}

class range {
  Integer start_, end_, step_;

public:
  struct range_iterator {
    Integer val, step_;
    range_iterator(Integer v, Integer step) : val(v), step_(step) {}
    Integer operator*() { return val; }
    void operator++() { val += step_; }
    bool operator!=(range_iterator &x) {
      return step_ > 0 ? val < x.val : val > x.val;
    }
  };
  range(Integer len) : start_(0), end_(len), step_(1) {}
  range(Integer start, Integer end) : start_(start), end_(end), step_(1) {}
  range(Integer start, Integer end, Integer step)
      : start_(start), end_(end), step_(step) {}
  range_iterator begin() { return range_iterator(start_, step_); }
  range_iterator end() { return range_iterator(end_, step_); }
};

inline string operator"" _s(const char *str, size_t size) {
  return move(string(str));
}
constexpr Integer my_pow(Integer x, Integer k, Integer z = 1) {
  return k == 0    ? z
         : k == 1  ? z * x
         : (k & 1) ? my_pow(x * x, k >> 1, z * x)
                   : my_pow(x * x, k >> 1, z);
}
constexpr Integer my_pow_mod(Integer x, Integer k, Integer M, Integer z = 1) {
  return k == 0    ? z % M
         : k == 1  ? z * x % M
         : (k & 1) ? my_pow_mod(x * x % M, k >> 1, M, z * x % M)
                   : my_pow_mod(x * x % M, k >> 1, M, z);
}
constexpr unsigned long long operator"" _ten(unsigned long long value) {
  return my_pow(10, value);
}

inline int k_bit(Integer x, int k) { return (x >> k) & 1; } // 0-indexed

mt19937 mt(chrono::duration_cast<chrono::nanoseconds>(
               chrono::steady_clock::now().time_since_epoch())
               .count());

template <class T> string join(const vector<T> &v, const string &sep) {
  stringstream ss;
  for (size_t i = 0; i < v.size(); i++) {
    if (i > 0)
      ss << sep;
    ss << v[i];
  }
  return ss.str();
}

inline string operator*(string s, int k) {
  string ret;
  while (k) {
    if (k & 1)
      ret += s;
    s += s;
    k >>= 1;
  }
  return ret;
}
} // namespace
constexpr long long mod = 9_ten + 7;

struct node {
  long long min = 0;
  long long tmp = 0;
};

node merge(node a, node b) {
  // assert(a.tmp == 0);
  // assert(b.tmp == 0);
  return node{
      min(a.min, b.min),
      0,
  };
}

struct node_wrap {
  node n;
  bool ok = false;
};

node_wrap merge(node_wrap a, node_wrap b) {
  if (a.ok && b.ok)
    return node_wrap{merge(a.n, b.n), true};
  a.n.tmp = 0;
  b.n.tmp = 0;
  if (a.ok)
    return a;
  if (b.ok)
    return b;
  return a;
}

struct seg_tree {
  struct node_info {
    int lb; // [lb, ub)
    int ub;
  };

  int n;
  vector<node_wrap> nodes;
  vector<node_info> info;

  inline int get_left(int id) { return id * 2 + 0; }
  inline int get_right(int id) { return id * 2 + 1; }

  inline bool has_intersection(int lb, int ub, int i) {
    return !(ub <= info[i].lb || info[i].ub <= lb);
  }

  inline bool is_included(int lb, int ub, int i) {
    return lb <= info[i].lb && info[i].ub <= ub;
  }

  void push(int i) {
    for (int k : {get_left(i), get_right(i)}) {
      if (nodes.size() <= k)
        continue;
      nodes[k].n.min += nodes[i].n.tmp;
      nodes[k].n.tmp += nodes[i].n.tmp;
    }
    nodes[i].n.tmp = 0;
  }

  // [lb,ub)
  node_wrap add(int lb, int ub, int x, int i = 1) {
    if (i >= nodes.size())
      return node_wrap();
    if (has_intersection(lb, ub, i) == false)
      return nodes[i];
    push(i);
    if (is_included(lb, ub, i)) {
      nodes[i].n.min += x;
      nodes[i].n.tmp += x;
      return nodes[i];
    }
    auto a = add(lb, ub, x, get_left(i));
    auto b = add(lb, ub, x, get_right(i));
    nodes[i] = merge(a, b);
    return nodes[i];
  }

  // [lb, ub)
  node_wrap get(int lb, int ub, int i = 1) {
    if (i >= nodes.size())
      return node_wrap();
    if (has_intersection(lb, ub, i) == false)
      return node_wrap();
    push(i);
    if (is_included(lb, ub, i)) {
      return nodes[i];
    }
    auto a = get(lb, ub, get_left(i));
    auto b = get(lb, ub, get_right(i));
    return merge(a, b);
  }

  void init(int n) {
    int m = 1;
    while (m <= n)
      m <<= 1;
    this->n = m;
    nodes = vector<node_wrap>(m + m, node_wrap{node(), true});
    info = vector<node_info>(m + m);
    for (int i = 0; i < m; i++) {
      info[i + m].lb = i;
      info[i + m].ub = i + 1;
    }
    for (int i = m - 1; i > 0; i--) {
      info[i].lb = info[get_left(i)].lb;
      info[i].ub = info[get_right(i)].ub;
    }
  }

  void set(int i, node x) {
    i += n;
    nodes[i].n = x;
    nodes[i].ok = true;
    for (i = i / 2; i > 0; i /= 2) {
      push(i);
      nodes[i] = merge(nodes[get_left(i)], nodes[get_right(i)]);
    }
  }
};

struct debug_tree {
  vector<int> val;
  void add(int lb, int ub, int x) {
    for (int i = lb; i < ub; i++) {
      val[i] += x;
    }
  }
  int get(int lb, int ub) {
    int res = -1 ^ (1 << 31);
    for (int i = lb; i < ub; i++) {
      res = min(res, val[i]);
    }
    return res;
  }
  void init(int n) { val.resize(n, 0); }
};

void solve() {
  int n;
  cin >> n;
  vector<int> k(n), l(n), r(n);
  for (int i : range(n)) {
    cin >> k[i], l[i], r[i];
    k[i]--;
    if (k[i] == n - 1)
      r[i] = 0;
  }

  vector<int> left, right;
  for (int i : range(n)) {
    if (l[i] >= r[i])
      left.push_back(i);
    else
      right.push_back(i);
  }

  sort(begin(left), end(left), [&](int i, int j) {
    if (l[i] - r[i] == l[j] - r[j]) {
      return k[i] > k[j];
    } else {
      return l[i] - r[i] > l[j] - r[j];
    }
  });

  sort(begin(right), end(right), [&](int i, int j) {
    if (r[i] - l[i] == r[j] - l[j]) {
      return k[i] < k[j];
    } else {
      return r[i] - l[i] > r[j] - l[j];
    }
  });

  vector<long long> dp_l(n + 1, 0);
  for (int i : left) {
    dp_l[0] += min(l[i], r[i]);
  }
  {
    seg_tree tree;
    // debug_tree tree;
    tree.init(n + 1);
    for (int i = 0; i < n; i++) {
      tree.set(i, node{0, 0});
    }
    int used = 0;
    for (int i : left) {
      int x = k[i];
      if (used <= x) {
        if (used < x)
          tree.add(used, x, 1);
        dp_l[used + 1] = dp_l[used] + l[i] - r[i];
        used++;
      } else {
        auto tmp = tree.get(x, used);
        // if(tmp == 0){
        if (tmp.n.min == 0) {
          continue;
        }
        tree.add(x, used, -1);
        dp_l[used + 1] = dp_l[used] + l[i] - r[i];
        used++;
      }
    }
  }

  vector<long long> dp_r(n + 1, 0);
  for (int i : right) {
    dp_r[0] += min(l[i], r[i]);
  }
  {
    for (auto &kk : k) {
      kk = kk + 1;
      kk = n - 1 - kk;
    }
    seg_tree tree;
    // debug_tree tree;
    tree.init(n + 1);
    for (int i = 0; i < n; i++) {
      tree.set(i, node{0, 0});
    }
    int used = 0;
    for (int i : right) {
      int x = k[i];
      if (used <= x) {
        if (used < x)
          tree.add(used, x, 1);
        dp_r[used + 1] = dp_r[used] + r[i] - l[i];
        used++;
      } else {
        auto tmp = tree.get(x, used);
        if (tmp.n.min == 0) {
          // if(tmp == 0){
          continue;
        }
        tree.add(x, used, -1);
        dp_r[used + 1] = dp_r[used] + r[i] - l[i];
        used++;
      }
    }
  }

  long long ans = 0;
  ans += *max_element(dp_l.begin(), dp_l.end());
  ans += *max_element(dp_r.begin(), dp_r.end());
  println(ans);
}

int main() {
  int n;
  cin >> n;
  while (n--) {
    solve();
  }

  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <cassert>
#include <chrono>
#include <cmath>
#include <cstdio>
#include <ctime>
#include <functional>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
using namespace std;

namespace {
using Integer = long long; //__int128;
template <class T, class S> istream &operator>>(istream &is, pair<T, S> &p) {
  return is >> p.first >> p.second;
}
template <class T> istream &operator>>(istream &is, vector<T> &vec) {
  for (T &val : vec)
    is >> val;
  return is;
}
template <class T> istream &operator,(istream &is, T &val) { return is >> val; }
template <class T, class S>
ostream &operator<<(ostream &os, const pair<T, S> &p) {
  return os << p.first << " " << p.second;
}
template <class T> ostream &operator<<(ostream &os, const vector<T> &vec) {
  for (size_t i = 0; i < vec.size(); i++)
    os << vec[i] << (i == vec.size() - 1 ? "" : " ");
  return os;
}
template <class T> ostream &operator,(ostream &os, const T &val) {
  return os << " " << val;
}

template <class H> void print(const H &head) { cout << head; }
template <class H, class... T> void print(const H &head, const T &...tail) {
  cout << head << " ";
  print(tail...);
}
template <class... T> void println(const T &...values) {
  print(values...);
  cout << endl;
}

template <class H> void eprint(const H &head) { cerr << head; }
template <class H, class... T> void eprint(const H &head, const T &...tail) {
  cerr << head << " ";
  eprint(tail...);
}
template <class... T> void eprintln(const T &...values) {
  eprint(values...);
  cerr << endl;
}

class range {
  Integer start_, end_, step_;

public:
  struct range_iterator {
    Integer val, step_;
    range_iterator(Integer v, Integer step) : val(v), step_(step) {}
    Integer operator*() { return val; }
    void operator++() { val += step_; }
    bool operator!=(range_iterator &x) {
      return step_ > 0 ? val < x.val : val > x.val;
    }
  };
  range(Integer len) : start_(0), end_(len), step_(1) {}
  range(Integer start, Integer end) : start_(start), end_(end), step_(1) {}
  range(Integer start, Integer end, Integer step)
      : start_(start), end_(end), step_(step) {}
  range_iterator begin() { return range_iterator(start_, step_); }
  range_iterator end() { return range_iterator(end_, step_); }
};

inline string operator"" _s(const char *str, size_t size) {
  return move(string(str));
}
constexpr Integer my_pow(Integer x, Integer k, Integer z = 1) {
  return k == 0    ? z
         : k == 1  ? z * x
         : (k & 1) ? my_pow(x * x, k >> 1, z * x)
                   : my_pow(x * x, k >> 1, z);
}
constexpr Integer my_pow_mod(Integer x, Integer k, Integer M, Integer z = 1) {
  return k == 0    ? z % M
         : k == 1  ? z * x % M
         : (k & 1) ? my_pow_mod(x * x % M, k >> 1, M, z * x % M)
                   : my_pow_mod(x * x % M, k >> 1, M, z);
}
constexpr unsigned long long operator"" _ten(unsigned long long value) {
  return my_pow(10, value);
}

inline int k_bit(Integer x, int k) { return (x >> k) & 1; } // 0-indexed

mt19937 mt(chrono::duration_cast<chrono::nanoseconds>(
               chrono::steady_clock::now().time_since_epoch())
               .count());

template <class T> string join(const vector<T> &v, const string &sep) {
  stringstream ss;
  for (size_t i = 0; i < v.size(); i++) {
    if (i > 0)
      ss << sep;
    ss << v[i];
  }
  return ss.str();
}

inline string operator*(string s, int k) {
  string ret;
  while (k) {
    if (k & 1)
      ret += s;
    s += s;
    k >>= 1;
  }
  return ret;
}
} // namespace
constexpr long long mod = 9_ten + 7;

struct node {
  long long min = 1ll << 58;
  long long tmp = 0;
};

node merge(node a, node b) {
  // assert(a.tmp == 0);
  // assert(b.tmp == 0);
  return node{
      min(a.min, b.min),
      0,
  };
}

struct node_wrap {
  node n;
  bool ok = false;
};

node_wrap merge(node_wrap a, node_wrap b) {
  if (a.ok && b.ok)
    return node_wrap{merge(a.n, b.n), true};
  a.n.tmp = 0;
  b.n.tmp = 0;
  if (a.ok)
    return a;
  if (b.ok)
    return b;
  return a;
}

struct seg_tree {
  int n;
  vector<node_wrap> nodes;

  void push(int i) {
    int l = i * 2;
    int r = i * 2 + 1;
    for (int k : {l, r}) {
      if (nodes.size() <= k)
        continue;
      if (nodes[k].ok == false) {
        nodes[k] = node_wrap{node{0, 0}, true};
      }
      nodes[k].n.min += nodes[i].n.tmp;
      nodes[k].n.tmp += nodes[i].n.tmp;
    }
    nodes[i].n.tmp = 0;
  }

  node_wrap add(int lb, int ub, int l, int u, int i, int x) {
    if (i >= nodes.size())
      return node_wrap{};
    if (u <= lb || ub <= l)
      return node_wrap{};
    push(i);
    if (lb <= l && u <= ub) {
      if (nodes[i].ok == false) {
        nodes[i] = node_wrap{node{0, 0}, true};
      }
      nodes[i].n.min += x;
      nodes[i].n.tmp += x;
      return nodes[i];
    }
    auto a = add(lb, ub, l, (l + u) / 2, i * 2, x);
    auto b = add(lb, ub, (l + u) / 2, u, i * 2 + 1, x);
    nodes[i] = merge(a, b);
    return nodes[i];
  }

  // [lb,ub)
  void add(int lb, int ub, int x) {
    int i = 1;
    int l = 0;
    int u = n;
    add(lb, ub, l, u, i, x);
  }

  node_wrap get(int lb, int ub, int l, int u, int i) {
    if (i >= nodes.size())
      return node_wrap{};
    if (u <= lb || ub <= l)
      return node_wrap{};
    push(i);
    if (lb <= l && u <= ub) {
      return nodes[i];
    }
    auto a = get(lb, ub, l, (l + u) / 2, i * 2);
    auto b = get(lb, ub, (l + u) / 2, u, i * 2 + 1);
    return merge(a, b);
  }

  // [lb, ub)
  node_wrap get(int lb, int ub) {
    int i = 1;
    int l = 0;
    int u = n;
    return get(lb, ub, l, u, i);
  }

  void init(int n) {
    int m = 1;
    while (m <= n)
      m <<= 1;
    this->n = m;
    nodes = vector<node_wrap>(m + m);
  }

  void set(int i, node x) {
    i += n;
    nodes[i].n = x;
    nodes[i].ok = true;
    for (i = i / 2; i > 0; i /= 2) {
      int l = i * 2;
      int r = l + 1;
      nodes[i] = merge(nodes[l], nodes[r]);
    }
  }
};

struct debug_tree {
  vector<int> val;
  void add(int lb, int ub, int x) {
    for (int i = lb; i < ub; i++) {
      val[i] += x;
    }
  }
  int get(int lb, int ub) {
    int res = -1 ^ (1 << 31);
    for (int i = lb; i < ub; i++) {
      res = min(res, val[i]);
    }
    return res;
  }
  void init(int n) { val.resize(n, 0); }
};

void solve() {
  int n;
  cin >> n;
  vector<int> k(n), l(n), r(n);
  for (int i : range(n)) {
    cin >> k[i], l[i], r[i];
    k[i]--;
    if (k[i] == n - 1)
      r[i] = 0;
  }

  vector<int> left, right;
  for (int i : range(n)) {
    if (l[i] >= r[i])
      left.push_back(i);
    else
      right.push_back(i);
  }

  sort(begin(left), end(left), [&](int i, int j) {
    if (l[i] - r[i] == l[j] - r[j]) {
      return k[i] > k[j];
    } else {
      return l[i] - r[i] > l[j] - r[j];
    }
  });

  sort(begin(right), end(right), [&](int i, int j) {
    if (r[i] - l[i] == r[j] - l[j]) {
      return k[i] < k[j];
    } else {
      return r[i] - l[i] > r[j] - l[j];
    }
  });

  vector<long long> dp_l(n + 1, 0);
  for (int i : left) {
    dp_l[0] += min(l[i], r[i]);
  }
  {
    // seg_tree tree;
    debug_tree tree;
    tree.init(n + 1);
    // for(int i=0; i<n; i++){
    //   tree.set(i, node{0, 0});
    // }
    int used = 0;
    for (int i : left) {
      int x = k[i];
      if (used <= x) {
        if (used < x)
          tree.add(used, x, 1);
        dp_l[used + 1] = dp_l[used] + l[i] - r[i];
        used++;
      } else {
        auto tmp = tree.get(x, used);
        if (tmp == 0) {
          // if(tmp.n.min == 0){
          continue;
        }
        tree.add(x, used, -1);
        dp_l[used + 1] = dp_l[used] + l[i] - r[i];
        used++;
      }
    }
  }

  vector<long long> dp_r(n + 1, 0);
  for (int i : right) {
    dp_r[0] += min(l[i], r[i]);
  }
  {
    for (auto &kk : k) {
      kk = kk + 1;
      kk = n - 1 - kk;
    }
    // seg_tree tree;
    debug_tree tree;
    tree.init(n + 1);
    // for(int i=0; i<n; i++){
    //   tree.set(i, node{0, 0});
    // }
    int used = 0;
    for (int i : right) {
      int x = k[i];
      if (used <= x) {
        if (used < x)
          tree.add(used, x, 1);
        dp_r[used + 1] = dp_r[used] + r[i] - l[i];
        used++;
      } else {
        auto tmp = tree.get(x, used);
        // if(tmp.n.min == 0){
        if (tmp == 0) {
          continue;
        }
        tree.add(x, used, -1);
        dp_r[used + 1] = dp_r[used] + r[i] - l[i];
        used++;
      }
    }
  }

  long long ans = 0;
  ans += *max_element(dp_l.begin(), dp_l.end());
  ans += *max_element(dp_r.begin(), dp_r.end());
  println(ans);
}

int main() {
  int n;
  cin >> n;
  while (n--) {
    solve();
  }

  return 0;
}