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

#include <algorithm>
#include <iostream>
#include <sstream>
#include <vector>
using namespace std;

long long GetNewSat(int d, int newT, const vector<int> &ts,
                    const vector<int> &cs, const vector<vector<int>> &ss,
                    long long sat, int D) {
  int prevT = ts[d];
  int prevTdaysBefore = distance(ts.rbegin() + (D - d),
                                 find(ts.rbegin() + (D - d), ts.rend(), prevT));
  int prevTdaysAfter =
      distance(ts.begin() + d, find(ts.begin() + d + 1, ts.end(), prevT));
  int newTdaysBefore = distance(ts.rbegin() + (D - d),
                                find(ts.rbegin() + (D - d), ts.rend(), newT));
  int newTdaysAfter =
      distance(ts.begin() + d, find(ts.begin() + d + 1, ts.end(), newT));
  sat -= cs[prevT] * (prevTdaysBefore + 1) * prevTdaysAfter;
  sat += cs[newT] * (newTdaysBefore + 1) * newTdaysAfter;
  sat -= ss[d][prevT];
  sat += ss[d][newT];
  return sat;
}

int main() {
  auto &in = cin;

  const int numA = 26;
  int D;
  in >> D;
  vector<int> cs;
  vector<vector<int>> ss;
  for (int i = 0; i < numA; ++i) {
    int c;
    in >> c;
    cs.push_back(c);
  }
  for (int i = 0; i < D; ++i) {
    vector<int> s;
    for (int j = 0; j < numA; ++j) {
      int is;
      in >> is;
      s.push_back(is);
    }
    ss.emplace_back(move(s));
  }

  long long sat = 0;
  vector<int> ds(numA);
  vector<int> ts;
  for (int i = 0; i < D; ++i) {
    int t = distance(ss[i].begin(), max_element(ss[i].begin(), ss[i].end()));
    ts.push_back(t);
    sat += ss[i][t];
    for (int j = 0; j < numA; ++j)
      ++ds[j];
    ds[t] = 0;
    for (int j = 0; j < numA; ++j)
      sat -= cs[j] * ds[j];
  }

  int M = 10000;
  int dif = 1;
  for (int i = 0; i < M; ++i) {
    for (int d = 0; d < D; ++d) {
      int newT = ts[d] + dif;
      dif = (dif + 1) % numA;
      newT %= numA;
      long long newSat = GetNewSat(d, newT, ts, cs, ss, sat, D);
      if (newSat > sat) {
        sat = newSat;
        ts[d] = newT;
      }
    }
  }

  for (int i = 0; i < D; ++i) {
    cout << ts[i] + 1 << "\\n";
  }

  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <iostream>
#include <sstream>
#include <vector>
using namespace std;

long long GetNewSat(int d, int newT, const vector<int> &ts,
                    const vector<int> &cs, const vector<vector<int>> &ss,
                    long long sat, int D) {
  int prevT = ts[d];
  int prevTdaysBefore = distance(ts.rbegin() + (D - d),
                                 find(ts.rbegin() + (D - d), ts.rend(), prevT));
  int prevTdaysAfter =
      distance(ts.begin() + d, find(ts.begin() + d + 1, ts.end(), prevT));
  int newTdaysBefore = distance(ts.rbegin() + (D - d),
                                find(ts.rbegin() + (D - d), ts.rend(), newT));
  int newTdaysAfter =
      distance(ts.begin() + d, find(ts.begin() + d + 1, ts.end(), newT));
  sat -= cs[prevT] * (prevTdaysBefore + 1) * prevTdaysAfter;
  sat += cs[newT] * (newTdaysBefore + 1) * newTdaysAfter;
  sat -= ss[d][prevT];
  sat += ss[d][newT];
  return sat;
}

int main() {
  auto &in = cin;

  const int numA = 26;
  int D;
  in >> D;
  vector<int> cs;
  vector<vector<int>> ss;
  for (int i = 0; i < numA; ++i) {
    int c;
    in >> c;
    cs.push_back(c);
  }
  for (int i = 0; i < D; ++i) {
    vector<int> s;
    for (int j = 0; j < numA; ++j) {
      int is;
      in >> is;
      s.push_back(is);
    }
    ss.emplace_back(move(s));
  }

  long long sat = 0;
  vector<int> ds(numA);
  vector<int> ts;
  for (int i = 0; i < D; ++i) {
    int t = distance(ss[i].begin(), max_element(ss[i].begin(), ss[i].end()));
    ts.push_back(t);
    sat += ss[i][t];
    for (int j = 0; j < numA; ++j)
      ++ds[j];
    ds[t] = 0;
    for (int j = 0; j < numA; ++j)
      sat -= cs[j] * ds[j];
  }

  int M = 70000;
  int dif = 1;
  for (int i = 0; i < M; ++i) {
    for (int d = 0; d < D; ++d) {
      int newT = ts[d] + dif;
      dif = (dif + 1) % numA;
      long long newSat = GetNewSat(d, newT, ts, cs, ss, sat, D);
      if (newSat > sat) {
        sat = newSat;
        ts[d] = newT;
      }
    }
  }

  for (int i = 0; i < D; ++i) {
    cout << ts[i] + 1 << "\\n";
  }

  return 0;
}