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


#include <algorithm>
#include <chrono>
#include <deque>
#include <fstream>
#include <functional>
#include <iostream>
#include <iterator>
#include <limits>
#include <map>
#include <math.h>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#define forr(i, start, count) for (int i = (start); i < (start) + (count); ++i)
#define set_map_includes(set, elt) (set.find((elt)) != set.end())
#define readint(i)                                                             \\
  int i;                                                                       \\
  cin >> i
#define readll(i)                                                              \\
  ll i;                                                                        \\
  cin >> i
#define readdouble(i)                                                          \\
  double i;                                                                    \\
  cin >> i
#define readstring(s)                                                          \\
  string s;                                                                    \\
  cin >> s

typedef long long ll;

using namespace std;

ll modd = 1000 * 1000 * 1000 + 7;

ll binary_search(function<bool(ll)> func, ll start, ll end) {
  /*      func:int ->bool
          returns smallest int x where func(x) evaluates to true, searches in
     [start,end), it is assumed the values are false, .. , false, true ...
           */
  if (end <= start) {
    return end;
  } // has to be here, otherwise func(end-1) in next line could be a problem
  if (!func(end - 1)) {
    return end;
  }
  while (end - start > 1) {
    ll mid = (start + end) / 2;
    if (func(mid)) {
      end = mid;
    } else {
      start = mid;
    }
  }
  if (func(start)) {
    return start;
  } else {
    return end;
  }
};

template <class T> class depth_first_search_iterator {
public:
  stack<pair<T, T>> todo;
  function<bool(T)> notvisited;
  function<void(T)> mark_visited;
  function<set<T>(T)> neighbors;

  depth_first_search_iterator() {}

  depth_first_search_iterator(T start, function<bool(T)> notv,
                              function<void(T)> mark_v,
                              function<set<T>(T)> neigh)
      : notvisited(notv), mark_visited(mark_v), neighbors(neigh), end_(false),
        current(start), prev(start) {
    todo.push(make_pair(start, start));
    operator++();
  }

  void operator++() {
    end_ = true;
    while (!todo.empty()) {
      current = todo.top().first;
      prev = todo.top().second;
      todo.pop();
      if (notvisited(current)) {
        mark_visited(current);
        for (const T x : neighbors(current)) {
          todo.push(make_pair(x, current));
        }
        end_ = false;
        break;
      }
    }
  }

  T operator*() { return current; }

  T previous() { return prev; }

  bool end() { return end_; }

private:
  T current; // contains current node
  T prev;    // contains previous node
  bool end_;
};

template <class T> class breadth_first_search_iterator {
public:
  queue<pair<T, T>> todo;
  function<bool(T)> notvisited;
  function<void(T)> mark_visited;
  function<set<T>(T)> neighbors;

  breadth_first_search_iterator() {}

  breadth_first_search_iterator(T start, function<bool(T)> notv,
                                function<void(T)> mark_v,
                                function<set<T>(T)> neigh)
      : notvisited(notv), mark_visited(mark_v), neighbors(neigh), end_(false),
        current(start), prev(start) {
    todo.push(make_pair(start, start));
    operator++();
  }

  void operator++() {
    end_ = true;
    while (!todo.empty()) {
      current = todo.front().first;
      prev = todo.front().second;
      todo.pop();
      if (notvisited(current)) {
        mark_visited(current);
        for (const T x : neighbors(current)) {
          todo.push(make_pair(x, current));
        }
        end_ = false;
        break;
      }
    }
  }

  T operator*() { return current; }

  T previous() { return prev; }

  bool end() { return end_; }

private:
  T current; // current node
  T prev;    // previous node
  bool end_;
};

template <class T> class DirectedGraph {
  // allows only single connection between two vertices
public:
  vector<T> vertices;
  vector<set<int>> neighbors;
  depth_first_search_iterator<int> dfs_iterator;
  breadth_first_search_iterator<int> bfs_iterator;
  vector<bool> visited;

  DirectedGraph() {}

  DirectedGraph(int n, T default_val = 0) {
    for (int i = 0; i < n; ++i) {
      AddVertex(default_val);
    }
  }

  bool EdgeExists(int i_from, int j_to) {
    return (neighbors[i_from].find(j_to) != neighbors[i_from].end());
  }

  void AddVertex(T val) {
    vertices.push_back(val);
    neighbors.push_back(set<int>());
  }

  void AddEdge(int i_from, int j_to) {
    if (!EdgeExists(i_from, j_to)) {
      neighbors[i_from].insert(j_to);
    }
  }

  T &operator[](int i) { return vertices[i]; }

  void RemoveEdge(int i_from, int j_to) {
    if (EdgeExists(i_from, j_to)) {
      neighbors[i_from].erase(
          find(neighbors[i_from].begin(), neighbors[i_from].end(), j_to));
    }
  }

  DirectedGraph<T> Transpose() {
    DirectedGraph<T> ret;
    for (auto x : vertices) {
      ret.AddVertex(x);
    }
    for (int i_from = 0; i_from < vertices.size(); ++i_from) {
      for (int j_to : neighbors[i_from]) {
        ret.AddEdge(j_to, i_from);
      }
    }
    return ret;
  }

  int vertices_count() { return vertices.size(); }

  void dfs_init(int start) {
    visited = vector<bool>(vertices.size(), false);
    dfs_iterator = depth_first_search_iterator<int>(
        start, [this](int x) { return !visited[x]; },
        [this](int x) { visited[x] = true; },
        [this](int x) { return neighbors[x]; });
  }

  void bfs_init(int start) {
    visited = vector<bool>(vertices.size(), false);
    bfs_iterator = breadth_first_search_iterator<int>(
        start, [this](int x) { return !visited[x]; },
        [this](int x) { visited[x] = true; },
        [this](int x) { return neighbors[x]; });
  }

  void print_() {
    forr(i, 0, vertices.size()) {
      for (auto x : neighbors[i]) {
        cout << i << "->" << x << endl;
      }
    }
  }
};

template <class T> struct SegmentData {
  int left, right;   // node represents interval [left, right)
  T max_in_interval; // func of all elements in this interval
  int prev;          // index of parent node -- equals -1 for root node
};

template <class T> class SegmentTree : public DirectedGraph<SegmentData<T>> {
  // creates segment tree, with func(kao max) values at nodes; edges go from
  // root down to children
public:
  int root;       // index of root
  int data_size;  // size of data array
  vector<T> data; // data

  SegmentTree(int n, function<T(T, T)> f, T zero_v)
      : default_val(zero_v), data(vector<T>(n, zero_v)), data_size(n),
        singleton_index(vector<int>(n, 0)), func(f) {
    // creates func segment tree
    root = HelperCreate(0, data.size());
    HelperSetPrev(root, -1);
  }

  T GetIntervalEvaluation(int x, int y) { // returns func of interval [x,y)
    if (x >= y) {
      return default_val;
    } // this is where you determine what should be returned for empty interval
    return GetMaxHelper(x, y, root);
  }

  void SetVal(int index, T val) {
    data[index] = val;
    int curr = singleton_index[index];
    this->vertices[curr].max_in_interval = val;
    while (this->vertices[curr].prev >= 0) {
      curr = this->vertices[curr].prev;
      T ret = this->vertices[*(this->neighbors[curr].begin())].max_in_interval;
      if (this->neighbors[curr].size() > 1) {
        if (this->vertices[*(this->neighbors[curr].begin())].left <
            this->vertices[*next(this->neighbors[curr].begin())].right) {
          ret = func(ret, this->vertices[*next(this->neighbors[curr].begin())]
                              .max_in_interval);
        } else {
          ret = func(this->vertices[*next(this->neighbors[curr].begin())]
                         .max_in_interval,
                     ret);
        }
      }
      this->vertices[curr].max_in_interval = ret;
    }
  }

  T operator[](int i) { return data[i]; }

private:
  T default_val; // zero/default value for data array
  function<T(T, T)>
      func; // node is func of both children -- func should be associative
  vector<int> singleton_index; // singleton_index[i] gives index of node in the
                               // graph for singleton [i,i+1)

  void HelperSetPrev(int root, int prev_) {
    this->vertices[root].prev = prev_;
    for (auto x : this->neighbors[root]) {
      HelperSetPrev(x, root);
    }
  }

  int HelperCreate(int left, int right) {
    if (right - left == 1) {
      this->AddVertex({left, right, default_val, 0});
      singleton_index[left] = this->vertices.size() - 1;
      return this->vertices.size() - 1;
    }
    int mid = (left + right) / 2;
    int l1 = HelperCreate(left, mid);
    int r1 = HelperCreate(mid, right);
    this->AddVertex({left, right, default_val, 0});
    this->AddEdge(this->vertices.size() - 1, l1);
    this->AddEdge(this->vertices.size() - 1, r1);

    return this->vertices.size() - 1;
  }

  T GetMaxHelper(int x, int y, int root) {
    if ((x == this->vertices[root].left) && (y == this->vertices[root].right)) {
      return this->vertices[root].max_in_interval;
    }
    int left = *(this->neighbors[root].begin());
    int right = *next(this->neighbors[root].begin());
    if (this->vertices[left].left > this->vertices[right].left) {
      swap(left, right);
    }
    int mid = this->vertices[left].right;
    if ((x < mid) && (y > mid)) {
      return func(GetMaxHelper(x, mid, left), GetMaxHelper(mid, y, right));
    }
    if ((x < mid) && (y <= mid)) {
      return GetMaxHelper(x, y, left);
    }
    if ((x >= mid) && (y > mid)) {
      return GetMaxHelper(x, y, right);
    }
  }
};

#define LSB(i)                                                                 \\
  ((i) & -(i)) // zeroes all the bits except the least significant one
// One based indexing is assumed
const int SIZE = 500005;
int A[SIZE + 1];

int sum(int i) // Returns the sum from index 1 to i
{
  int sum = 0;
  while (i > 0)
    sum += A[i], i -= LSB(i);
  return sum;
}

void add(int i, int k) // Adds k to element with index i
{
  while (i <= SIZE)
    A[i] += k, i += LSB(i);
}

int main() {

  ios_base::sync_with_stdio(false);

  cout.precision(17);
  mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
  uniform_int_distribution<int> rand_gen(
      0, modd); // rand_gen(rng) gets the rand no

  //    auto start = chrono::steady_clock::now();

  //    readint(test_cases);
  int test_cases = 1;
  forr(t, 1, test_cases) {
    readint(n);
    readint(q);
    vector<int> c(n + 2);
    vector<int> colors(n + 1, 0);
    forr(i, 0, n) {
      readint(cc);
      c[i + 1] = colors[cc];
      colors[cc] = i + 1;
    }

    SegmentTree<int> seg1(
        c.size(), [](int x, int y) { return x + y; }, 0);

    vector<vector<pair<int, int>>> _end(n + 2);
    vector<pair<int, int>> query;
    forr(i, 0, q) {
      readint(l);
      readint(r);
      _end[l].push_back(make_pair(i, 0));
      _end[r + 1].push_back(make_pair(i, 1));
      query.push_back(make_pair(l, r));
    }
    //        cerr << "5" << endl;
    vector<int> ans(q);

    for (int i = _end.size() - 1; i >= 1; --i) {
      //            seg1.SetVal(c[i], seg1[c[i]]+1);
      add(c[i] + 1, 1);
      for (auto x : _end[i]) {
        if (x.second == 0) {
          //                ans[x.first] += seg1.GetIntervalEvaluation(0,
          //                query[x.first].first);
          ans[x.first] += sum(query[x.first].first);
        } else {
          //                ans[x.first] -= seg1.GetIntervalEvaluation(0,
          //                query[x.first].first);
          ans[x.first] -= sum(query[x.first].first);
        }
      }
    }
    //        cerr << "6" << endl;

    for (auto x : ans) {
      cout << x << endl;
    }
  }

  //    auto end = chrono::steady_clock::now();
  //    cerr << chrono::duration_cast<chrono::milliseconds>(end - start).count()
  //    << endl;

  return 0;
}

Error: TLE

Fixed submission:


#include <algorithm>
#include <chrono>
#include <deque>
#include <fstream>
#include <functional>
#include <iostream>
#include <iterator>
#include <limits>
#include <map>
#include <math.h>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#define forr(i, start, count) for (int i = (start); i < (start) + (count); ++i)
#define set_map_includes(set, elt) (set.find((elt)) != set.end())
#define readint(i)                                                             \\
  int i;                                                                       \\
  cin >> i
#define readll(i)                                                              \\
  ll i;                                                                        \\
  cin >> i
#define readdouble(i)                                                          \\
  double i;                                                                    \\
  cin >> i
#define readstring(s)                                                          \\
  string s;                                                                    \\
  cin >> s

typedef long long ll;

using namespace std;

ll modd = 1000 * 1000 * 1000 + 7;

ll binary_search(function<bool(ll)> func, ll start, ll end) {
  /*      func:int ->bool
          returns smallest int x where func(x) evaluates to true, searches in
     [start,end), it is assumed the values are false, .. , false, true ...
           */
  if (end <= start) {
    return end;
  } // has to be here, otherwise func(end-1) in next line could be a problem
  if (!func(end - 1)) {
    return end;
  }
  while (end - start > 1) {
    ll mid = (start + end) / 2;
    if (func(mid)) {
      end = mid;
    } else {
      start = mid;
    }
  }
  if (func(start)) {
    return start;
  } else {
    return end;
  }
};

template <class T> class depth_first_search_iterator {
public:
  stack<pair<T, T>> todo;
  function<bool(T)> notvisited;
  function<void(T)> mark_visited;
  function<set<T>(T)> neighbors;

  depth_first_search_iterator() {}

  depth_first_search_iterator(T start, function<bool(T)> notv,
                              function<void(T)> mark_v,
                              function<set<T>(T)> neigh)
      : notvisited(notv), mark_visited(mark_v), neighbors(neigh), end_(false),
        current(start), prev(start) {
    todo.push(make_pair(start, start));
    operator++();
  }

  void operator++() {
    end_ = true;
    while (!todo.empty()) {
      current = todo.top().first;
      prev = todo.top().second;
      todo.pop();
      if (notvisited(current)) {
        mark_visited(current);
        for (const T x : neighbors(current)) {
          todo.push(make_pair(x, current));
        }
        end_ = false;
        break;
      }
    }
  }

  T operator*() { return current; }

  T previous() { return prev; }

  bool end() { return end_; }

private:
  T current; // contains current node
  T prev;    // contains previous node
  bool end_;
};

template <class T> class breadth_first_search_iterator {
public:
  queue<pair<T, T>> todo;
  function<bool(T)> notvisited;
  function<void(T)> mark_visited;
  function<set<T>(T)> neighbors;

  breadth_first_search_iterator() {}

  breadth_first_search_iterator(T start, function<bool(T)> notv,
                                function<void(T)> mark_v,
                                function<set<T>(T)> neigh)
      : notvisited(notv), mark_visited(mark_v), neighbors(neigh), end_(false),
        current(start), prev(start) {
    todo.push(make_pair(start, start));
    operator++();
  }

  void operator++() {
    end_ = true;
    while (!todo.empty()) {
      current = todo.front().first;
      prev = todo.front().second;
      todo.pop();
      if (notvisited(current)) {
        mark_visited(current);
        for (const T x : neighbors(current)) {
          todo.push(make_pair(x, current));
        }
        end_ = false;
        break;
      }
    }
  }

  T operator*() { return current; }

  T previous() { return prev; }

  bool end() { return end_; }

private:
  T current; // current node
  T prev;    // previous node
  bool end_;
};

template <class T> class DirectedGraph {
  // allows only single connection between two vertices
public:
  vector<T> vertices;
  vector<set<int>> neighbors;
  depth_first_search_iterator<int> dfs_iterator;
  breadth_first_search_iterator<int> bfs_iterator;
  vector<bool> visited;

  DirectedGraph() {}

  DirectedGraph(int n, T default_val = 0) {
    for (int i = 0; i < n; ++i) {
      AddVertex(default_val);
    }
  }

  bool EdgeExists(int i_from, int j_to) {
    return (neighbors[i_from].find(j_to) != neighbors[i_from].end());
  }

  void AddVertex(T val) {
    vertices.push_back(val);
    neighbors.push_back(set<int>());
  }

  void AddEdge(int i_from, int j_to) {
    if (!EdgeExists(i_from, j_to)) {
      neighbors[i_from].insert(j_to);
    }
  }

  T &operator[](int i) { return vertices[i]; }

  void RemoveEdge(int i_from, int j_to) {
    if (EdgeExists(i_from, j_to)) {
      neighbors[i_from].erase(
          find(neighbors[i_from].begin(), neighbors[i_from].end(), j_to));
    }
  }

  DirectedGraph<T> Transpose() {
    DirectedGraph<T> ret;
    for (auto x : vertices) {
      ret.AddVertex(x);
    }
    for (int i_from = 0; i_from < vertices.size(); ++i_from) {
      for (int j_to : neighbors[i_from]) {
        ret.AddEdge(j_to, i_from);
      }
    }
    return ret;
  }

  int vertices_count() { return vertices.size(); }

  void dfs_init(int start) {
    visited = vector<bool>(vertices.size(), false);
    dfs_iterator = depth_first_search_iterator<int>(
        start, [this](int x) { return !visited[x]; },
        [this](int x) { visited[x] = true; },
        [this](int x) { return neighbors[x]; });
  }

  void bfs_init(int start) {
    visited = vector<bool>(vertices.size(), false);
    bfs_iterator = breadth_first_search_iterator<int>(
        start, [this](int x) { return !visited[x]; },
        [this](int x) { visited[x] = true; },
        [this](int x) { return neighbors[x]; });
  }

  void print_() {
    forr(i, 0, vertices.size()) {
      for (auto x : neighbors[i]) {
        cout << i << "->" << x << endl;
      }
    }
  }
};

template <class T> struct SegmentData {
  int left, right;   // node represents interval [left, right)
  T max_in_interval; // func of all elements in this interval
  int prev;          // index of parent node -- equals -1 for root node
};

template <class T> class SegmentTree : public DirectedGraph<SegmentData<T>> {
  // creates segment tree, with func(kao max) values at nodes; edges go from
  // root down to children
public:
  int root;       // index of root
  int data_size;  // size of data array
  vector<T> data; // data

  SegmentTree(int n, function<T(T, T)> f, T zero_v)
      : default_val(zero_v), data(vector<T>(n, zero_v)), data_size(n),
        singleton_index(vector<int>(n, 0)), func(f) {
    // creates func segment tree
    root = HelperCreate(0, data.size());
    HelperSetPrev(root, -1);
  }

  T GetIntervalEvaluation(int x, int y) { // returns func of interval [x,y)
    if (x >= y) {
      return default_val;
    } // this is where you determine what should be returned for empty interval
    return GetMaxHelper(x, y, root);
  }

  void SetVal(int index, T val) {
    data[index] = val;
    int curr = singleton_index[index];
    this->vertices[curr].max_in_interval = val;
    while (this->vertices[curr].prev >= 0) {
      curr = this->vertices[curr].prev;
      T ret = this->vertices[*(this->neighbors[curr].begin())].max_in_interval;
      if (this->neighbors[curr].size() > 1) {
        if (this->vertices[*(this->neighbors[curr].begin())].left <
            this->vertices[*next(this->neighbors[curr].begin())].right) {
          ret = func(ret, this->vertices[*next(this->neighbors[curr].begin())]
                              .max_in_interval);
        } else {
          ret = func(this->vertices[*next(this->neighbors[curr].begin())]
                         .max_in_interval,
                     ret);
        }
      }
      this->vertices[curr].max_in_interval = ret;
    }
  }

  T operator[](int i) { return data[i]; }

private:
  T default_val; // zero/default value for data array
  function<T(T, T)>
      func; // node is func of both children -- func should be associative
  vector<int> singleton_index; // singleton_index[i] gives index of node in the
                               // graph for singleton [i,i+1)

  void HelperSetPrev(int root, int prev_) {
    this->vertices[root].prev = prev_;
    for (auto x : this->neighbors[root]) {
      HelperSetPrev(x, root);
    }
  }

  int HelperCreate(int left, int right) {
    if (right - left == 1) {
      this->AddVertex({left, right, default_val, 0});
      singleton_index[left] = this->vertices.size() - 1;
      return this->vertices.size() - 1;
    }
    int mid = (left + right) / 2;
    int l1 = HelperCreate(left, mid);
    int r1 = HelperCreate(mid, right);
    this->AddVertex({left, right, default_val, 0});
    this->AddEdge(this->vertices.size() - 1, l1);
    this->AddEdge(this->vertices.size() - 1, r1);

    return this->vertices.size() - 1;
  }

  T GetMaxHelper(int x, int y, int root) {
    if ((x == this->vertices[root].left) && (y == this->vertices[root].right)) {
      return this->vertices[root].max_in_interval;
    }
    int left = *(this->neighbors[root].begin());
    int right = *next(this->neighbors[root].begin());
    if (this->vertices[left].left > this->vertices[right].left) {
      swap(left, right);
    }
    int mid = this->vertices[left].right;
    if ((x < mid) && (y > mid)) {
      return func(GetMaxHelper(x, mid, left), GetMaxHelper(mid, y, right));
    }
    if ((x < mid) && (y <= mid)) {
      return GetMaxHelper(x, y, left);
    }
    if ((x >= mid) && (y > mid)) {
      return GetMaxHelper(x, y, right);
    }
  }
};

int main() {

  ios_base::sync_with_stdio(false);

  cout.precision(17);
  mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
  uniform_int_distribution<int> rand_gen(
      0, modd); // rand_gen(rng) gets the rand no

  //    auto start = chrono::steady_clock::now();

  //    readint(test_cases);
  int test_cases = 1;
  forr(t, 1, test_cases) {
    readint(n);
    readint(q);
    vector<int> c(n + 1);
    vector<int> colors(n + 1, 0);
    forr(i, 0, n) {
      readint(cc);
      c[i + 1] = colors[cc];
      colors[cc] = i + 1;
    }

    SegmentTree<int> seg1(
        c.size(), [](int x, int y) { return x + y; }, 0);
    //        cerr << "4" << endl;
    //        forr(i,0,c.size()) {
    //            seg1.SetVal(c[i], seg1[c[i]]+1);
    //        }

    vector<vector<pair<int, int>>> _end(n + 2);
    vector<pair<int, int>> query;
    forr(i, 0, q) {
      readint(l);
      readint(r);
      _end[l].push_back(make_pair(i, 0));
      _end[r + 1].push_back(make_pair(i, 1));
      query.push_back(make_pair(l, r));
    }
    //        cerr << "5" << endl;
    vector<int> ans(q);

    for (int i = _end.size() - 1; i >= 1; --i) {
      seg1.SetVal(c[i], seg1[c[i]] + 1);
      for (auto x : _end[i]) {
        if (x.second == 0) {
          ans[x.first] += seg1.GetIntervalEvaluation(0, query[x.first].first);
        } else {
          ans[x.first] -= seg1.GetIntervalEvaluation(0, query[x.first].first);
        }
      }
    }
    //        cerr << "6" << endl;

    for (auto x : ans) {
      cout << x << endl;
    }
  }

  //    auto end = chrono::steady_clock::now();
  //    cerr << chrono::duration_cast<chrono::milliseconds>(end - start).count()
  //    << endl;

  return 0;
}