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

#include <algorithm>
#include <chrono>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <list>
#include <map>
#include <queue>
#include <sstream>
#include <stack>
#include <string>
#include <vector>
#define ll long long

using namespace std;

long long MOD = 1000000007;

struct BIT {
  int vec[513];

  void insert(int x, int a) {
    while (x < 513) {
      // cout << x << endl;
      vec[x] += a;
      x += x & (-x);
    }
  }

  int get(int x) {
    int ret = 0;
    while (x > 0) {
      // cout << x << endl;
      ret += vec[x];
      x -= x & (-x);
    }
    return ret;
  }
};

int bin_search1(int l, int r, int a, BIT &bt) {
  int mid = (l + r) / 2;
  // cout << l << ' ' << mid << ' ' << r << endl;
  if (l == mid) {
    return mid;
  }

  if (bt.get(mid) > a) {
    return bin_search1(l, mid, a, bt);
  } else {
    return bin_search1(mid, r, a, bt);
  }
}

int bin_search2(int l, int r, int a, BIT &bt) {
  int mid = (l + r) / 2 + (abs(l) + r) % 2;
  // cout << l << ' ' << mid << ' ' << r << endl;
  if (r - l == 1) {
    return r;
  }

  if (bt.get(mid) < a) {
    return bin_search2(mid, r, a, bt);
  } else {
    return bin_search2(l, mid, a, bt);
  }
}

void cout_vec(int (&vec)[513]) {
  for (int i = 0; i < 513; i++) {
    if (i != 0) {
      cout << ' ';
    }
    cout << vec[i];
  }
  cout << endl;
}

int main() {
  chrono::system_clock::time_point start, end;
  start = chrono::system_clock::now();

  int D;
  cin >> D;
  vector<ll> c(26, 0);
  ll csum = 0;
  for (auto &x : c) {
    cin >> x;
    csum += x;
  }

  vector<vector<ll>> s(D, vector<ll>(26, 0));
  vector<ll> sum_(26, 0);
  for (int i = 0; i < D; i++) {
    for (int j = 0; j < 26; j++) {
      cin >> s[i][j];
      sum_[j] += s[i][j];
    }
  }

  int init = 0;
  ll max_ = 0;
  for (int i = 0; i < 26; i++) {
    if (max_ < sum_[i] + (c[i] - csum) * D * (D + 1) / 2) {
      max_ = sum_[i] + (c[i] - csum) * D * (D + 1) / 2;
      init = i;
    }
  }

  vector<int> ans(D + 1, init);
  vector<BIT> last_day(26);
  for (int i = 1; i <= D; i++) {
    last_day[init].insert(i, 1);
  }

  /*
  for(int i=1; i<=D; i++){
      cout << last_day[init].get(i) << endl;
  }
  */

  double time = 0.0;
  ll tmp = max_;
  while (time < 1900.) {
    // cout << time << endl;
    int d = rand() % D + 1;
    int x = rand() % 26;
    if (ans[d] == x) {
      // cout << time << endl;
      end = chrono::system_clock::now();
      time = static_cast<double>(
          chrono::duration_cast<chrono::microseconds>(end - start).count() /
          1000.0);
      continue;
    }

    int nd =
        bin_search1(d, D + 1, last_day[ans[d]].get(d), last_day[ans[d]]) + 1;
    int ld = bin_search2(-1, d, last_day[ans[d]].get(d) - 1, last_day[ans[d]]);
    int xnd = bin_search1(d, D + 1, last_day[x].get(d), last_day[x]) + 1;
    int xld = bin_search2(-1, d, last_day[x].get(d) - 1, last_day[x]);
    ll diff =
        c[ans[d]] * ((d - ld) * (d - ld + 1) / 2 + (nd - d) * (nd - d + 1) / 2 -
                     (nd - ld) * (nd - ld + 1) / 2) +
        c[x] * ((xnd - xld) * (xnd - xld + 1) / 2 -
                (d - xld) * (d - xld + 1) / 2 - (xnd - d) * (xnd - d + 1) / 2);

    // cout << x << ' ' << ld << ' ' << d << ' ' << nd << ' ' << xld << ' ' <<
    // xnd << endl; cout_vec(last_day[ans[d]].vec);
    if (diff > 0) {
      max_ = tmp;
      last_day[ans[d]].insert(d, -1);
      last_day[x].insert(d, 1);
      ans[d] = x;
    }
    end = chrono::system_clock::now();
    time = static_cast<double>(
        chrono::duration_cast<chrono::microseconds>(end - start).count() /
        1000.0);
  }
  for (int i = 1; i <= D; i++) {
    cout << ans[i] + 1 << '\\n';
  }
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <chrono>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <list>
#include <map>
#include <queue>
#include <sstream>
#include <stack>
#include <string>
#include <vector>
#define ll long long

using namespace std;

long long MOD = 1000000007;

struct BIT {
  int vec[513];

  void insert(int x, int a) {
    while (x < 513) {
      // cout << x << endl;
      vec[x] += a;
      x += x & (-x);
    }
  }

  int get(int x) {
    int ret = 0;
    while (x > 0) {
      // cout << x << endl;
      ret += vec[x];
      x -= x & (-x);
    }
    return ret;
  }
};

int bin_search1(int l, int r, int a, BIT &bt) {
  int mid = (l + r) / 2;
  // cout << l << ' ' << mid << ' ' << r << endl;
  if (l == mid) {
    return mid;
  }

  if (bt.get(mid) > a) {
    return bin_search1(l, mid, a, bt);
  } else {
    return bin_search1(mid, r, a, bt);
  }
}

int bin_search2(int l, int r, int a, BIT &bt) {
  int mid = (l + r) / 2 + (abs(l) + r) % 2;
  // cout << l << ' ' << mid << ' ' << r << endl;
  if (r - l == 1) {
    return r;
  }

  if (bt.get(mid) < a) {
    return bin_search2(mid, l, a, bt);
  } else {
    return bin_search2(l, mid, a, bt);
  }
}

void cout_vec(int (&vec)[513]) {
  for (int i = 0; i < 513; i++) {
    if (i != 0) {
      cout << ' ';
    }
    cout << vec[i];
  }
  cout << endl;
}

int main() {
  chrono::system_clock::time_point start, end;
  start = chrono::system_clock::now();

  int D;
  cin >> D;
  vector<ll> c(26, 0);
  ll csum = 0;
  for (auto &x : c) {
    cin >> x;
    csum += x;
  }

  vector<vector<ll>> s(D, vector<ll>(26, 0));
  vector<ll> sum_(26, 0);
  for (int i = 0; i < D; i++) {
    for (int j = 0; j < 26; j++) {
      cin >> s[i][j];
      sum_[j] += s[i][j];
    }
  }

  int init = 0;
  ll max_ = 0;
  for (int i = 0; i < 26; i++) {
    if (max_ < sum_[i] + (c[i] - csum) * D * (D + 1) / 2) {
      max_ = sum_[i] + (c[i] - csum) * D * (D + 1) / 2;
      init = i;
    }
  }

  vector<int> ans(D + 1, init);
  vector<BIT> last_day(26);
  for (int i = 1; i <= D; i++) {
    last_day[init].insert(i, 1);
  }

  /*
  for(int i=1; i<=D; i++){
      cout << last_day[init].get(i) << endl;
  }
  */

  double time = 0.0;
  ll tmp = max_;
  while (time < 1900.) {
    // cout << time << endl;
    int d = rand() % D + 1;
    int x = rand() % 26;
    if (ans[d] == x) {
      // cout << time << endl;
      end = chrono::system_clock::now();
      time = static_cast<double>(
          chrono::duration_cast<chrono::microseconds>(end - start).count() /
          1000.0);
      continue;
    }

    int nd =
        bin_search1(d, D + 1, last_day[ans[d]].get(d), last_day[ans[d]]) + 1;
    int ld = bin_search2(-1, d, last_day[ans[d]].get(d) - 1, last_day[ans[d]]);
    int xnd = bin_search1(d, D + 1, last_day[x].get(d), last_day[x]) + 1;
    int xld = bin_search2(-1, d, last_day[x].get(d) - 1, last_day[x]);
    tmp = max_ + c[ans[d]] * ((d - ld) * (d - ld + 1) / 2 +
                              (nd - d) * (nd - d + 1) / 2 -
                              (nd - ld) * (nd - ld + 1) / 2);
    tmp +=
        -c[x] * ((xnd - xld) * (xnd - xld + 1) / 2 +
                 (d - xld) * (d - xld + 1) / 2 + (xnd - d) * (xnd - d + 1) / 2);

    // cout << x << ' ' << ld << ' ' << d << ' ' << nd << ' ' <<
    // last_day[ans[d]].get(d) << ' ' << last_day[ans[d]].get(d+1) << endl;
    // cout_vec(last_day[ans[d]].vec);
    if (max_ < tmp) {
      max_ = tmp;
      last_day[ans[d]].insert(d, -1);
      last_day[x].insert(d, 1);
      ans[d] = x;
    }
    end = chrono::system_clock::now();
    time = static_cast<double>(
        chrono::duration_cast<chrono::microseconds>(end - start).count() /
        1000.0);
  }
  for (int i = 1; i <= D; i++) {
    cout << ans[i] + 1 << '\\n';
  }
}